This documentation is automatically generated by online-judge-tools/verification-helper
spec/test_runner.cr
- View this file on GitHub
- Last update: 2023-11-28 17:27:51+09:00
Depends on
- spec/convolution_spec.cr
- spec/dsu_spec.cr
- spec/fenwick_tree_spec.cr
- spec/lazy_seg_tree_spec.cr
- spec/math_spec.cr
- spec/min_cost_flow_spec.cr
- spec/mod_int_spec.cr
- spec/prime_spec.cr
- spec/priority_queue_spec.cr
- spec/red_black_tree_spec.cr
- spec/seg_tree_spec.cr
- spec/skew_heap_spec.cr
- spec/spec_helper.cr
- spec/string_spec.cr
Code
require "./*"# require "./*"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/convolution.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements [atcoder::convolution](https://atcoder.github.io/ac-library/master/document_en/convolution.html).
#
# ```
# a = [AtCoder::ModInt998244353.new(1_i64)] * 3
# AtCoder::Convolution.convolution(a, a) # => [1, 2, 3, 2, 1]
# ```
module Convolution
private def self.bit_reverse(n : Int, bit_length : Int)
ret = n.class.zero
bit_length.times do |i|
ret |= ((n >> i) & 1) << (bit_length - i - 1)
end
ret
end
# In-place execution of FFT operation.
# `T` must implement ModInt operations.
# Length of `a` must be power of 2.
private def self.fft(a : Array(T), g : T) forall T
size = a.size
bit_length = size.trailing_zeros_count
size.times do |i|
j = bit_reverse(i, bit_length)
if i < j
a[i], a[j] = a[j], a[i]
end
end
bit_length.times do |bit|
block_size = 1_i64 << bit
# Butterfly operation
block_size.times do |i|
w = (g ** (size // block_size // 2)) ** i
(size // (block_size * 2)).times do |j|
index1 = j * 2 * block_size + i
index2 = (j * 2 + 1) * block_size + i
a[index1], a[index2] = {
a[index1] + a[index2] * w,
a[index1] - a[index2] * w,
}
end
end
end
end
private def self.ifft(a, g)
fft(a, g.inv)
a.size.times do |i|
a[i] //= a.size
end
end
# Implements atcoder::convolution.convolution.
# TODO: Support for int
def self.convolution(a : Array(T), b : Array(T)) forall T
return [] of T if a.empty? || b.empty?
modulo = T::MOD
n = modulo - 1
result_size = a.size + b.size - 1
c = 1_i64 << n.trailing_zeros_count
if result_size > c
raise ArgumentError.new("With modulo #{modulo}, total array length must be less than or equal to #{c}")
end
fft_size = 1_i64
until fft_size >= result_size
fft_size <<= 1
end
input_a = Array(T).new(fft_size) { |i| i < a.size ? a[i] : T.zero }
input_b = Array(T).new(fft_size) { |i| i < b.size ? b[i] : T.zero }
primitive_root = AtCoder::Math.get_primitive_root(modulo)
g = T.new(primitive_root) ** (n // fft_size)
fft(input_a, g)
fft(input_b, g)
fft_size.times do |i|
input_a[i] *= input_b[i]
end
ifft(input_a, g)
input_a[0...result_size]
end
# Implements atcoder::convolution.convolution_ll.
def self.convolution_ll(a : Array(Int64), b : Array(Int64))
return [] of Int64 if a.empty? || b.empty?
a1 = a.map { |n| AtCoder::ModInt754974721.new(n) }
a2 = a.map { |n| AtCoder::ModInt167772161.new(n) }
a3 = a.map { |n| AtCoder::ModInt469762049.new(n) }
b1 = b.map { |n| AtCoder::ModInt754974721.new(n) }
b2 = b.map { |n| AtCoder::ModInt167772161.new(n) }
b3 = b.map { |n| AtCoder::ModInt469762049.new(n) }
c1 = convolution(a1, b1)
c2 = convolution(a2, b2)
c3 = convolution(a3, b3)
m1 = 754_974_721_i64
m2 = 167_772_161_i64
m3 = 469_762_049_i64
c1.zip(c2, c3).map do |n1, n2, n3|
p = AtCoder::Math.inv_mod(m1, m2)
tmp = (n2.val - n1.val) * p % m2
answer = n1.val + m1 * tmp
p = AtCoder::Math.inv_mod(m1 * m2, m3)
tmp = (((n3.val - answer) % m3) * p) % m3
answer + m1 * m2 * tmp
end
end
end
end
# require "../src/mod_int.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./math.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./prime.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./math.cr"
module AtCoder
# Implements [Ruby's Prime library](https://ruby-doc.com/stdlib/libdoc/prime/rdoc/Prime.html).
#
# ```
# AtCoder::Prime.first(7) # => [2, 3, 5, 7, 11, 13, 17]
# ```
module Prime
extend self
include Enumerable(Int64)
@@primes = [
2_i64, 3_i64, 5_i64, 7_i64, 11_i64, 13_i64, 17_i64, 19_i64,
23_i64, 29_i64, 31_i64, 37_i64, 41_i64, 43_i64, 47_i64,
53_i64, 59_i64, 61_i64, 67_i64, 71_i64, 73_i64, 79_i64,
83_i64, 89_i64, 97_i64, 101_i64,
]
def each(&)
index = 0
loop do
yield get_nth_prime(index)
index += 1
end
end
def prime_division(value : Int)
raise DivisionByZeroError.new if value == 0
int = typeof(value)
factors = [] of Tuple(typeof(value), typeof(value))
if value < 0
value = value.abs
factors << {int.new(-1), int.new(1)}
end
until prime?(value) || value == 1
factor = value
until prime?(factor)
factor = find_factor(factor)
end
count = 0
while value % factor == 0
value //= factor
count += 1
end
factors << {int.new(factor), int.new(count)}
end
if value > 1
factors << {value, int.new(1)}
end
factors.sort_by! { |(factor, _)| factor }
end
private def find_factor(n : Int)
# Factor of even numbers cannot be discovered by Pollard's Rho with f(x) = x^x+i
if n.even?
typeof(n).new(2)
else
pollard_rho(n).not_nil!
end
end
# Get single factor by Pollard's Rho Algorithm
private def pollard_rho(n : Int)
typeof(n).new(1).upto(n) do |i|
x = i
y = pollard_random_f(x, n, i)
loop do
x = pollard_random_f(x, n, i)
y = pollard_random_f(pollard_random_f(y, n, i), n, i)
gcd = (x - y).gcd(n)
if gcd == n
break
end
if gcd != 1
return gcd
end
end
end
end
@[AlwaysInline]
private def pollard_random_f(n : Int, mod : Int, seed : Int)
(AtCoder::Math.mul_mod(n, n, mod) + seed) % mod
end
private def extract_prime_division_base(prime_divisions_class : Array({T, T}).class) forall T
T
end
def int_from_prime_division(prime_divisions : Array({Int, Int}))
int_class = extract_prime_division_base(prime_divisions.class)
prime_divisions.reduce(int_class.new(1)) { |i, (factor, exponent)| i * factor ** exponent }
end
def prime?(value : Int)
# Obvious patterns
return false if value < 2
return true if value <= 3
return false if value.even?
return true if value < 9
if value < 0xffff
return false unless typeof(value).new(30).gcd(value % 30) == 1
7.step(by: 30, to: value) do |base|
break if base * base > value
if {0, 4, 6, 10, 12, 16, 22, 24}.any? { |i| value % (base + i) == 0 }
return false
end
end
return true
end
miller_rabin(value.to_i64)
end
private def miller_rabin(value)
d = value - 1
s = 0_i64
until d.odd?
d >>= 1
s += 1
end
miller_rabin_bases(value).each do |base|
next if base == value
x = AtCoder::Math.pow_mod(base.to_i64, d, value)
next if x == 1 || x == value - 1
is_composite = s.times.all? do
x = AtCoder::Math.mul_mod(x, x, value)
x != value - 1
end
return false if is_composite
end
true
end
# We can reduce time complexity of Miller-Rabin tests by testing against
# predefined bases which is enough to test against primarity in the given range.
# https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test
# ameba:disable Metrics/CyclomaticComplexity
private def miller_rabin_bases(value)
case
when value < 1_373_653_i64
[2, 3]
when value < 9_080_191_i64
[31, 73]
when value < 25_326_001_i64
[2, 3, 5]
when value < 3_215_031_751_i64
[2, 3, 5, 7]
when value < 4_759_123_141_i64
[2, 7, 61]
when value < 1_122_004_669_633_i64
[2, 13, 23, 1662803]
when value < 2_152_302_898_747_i64
[2, 3, 5, 7, 11]
when value < 3_474_749_660_383_i64
[2, 3, 5, 7, 11, 13]
when value < 341_550_071_728_321_i64
[2, 3, 5, 7, 11, 13, 17]
when value < 3_825_123_056_546_413_051_i64
[2, 3, 5, 7, 11, 13, 17, 19, 23]
else
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]
end
end
private def get_nth_prime(n)
while @@primes.size <= n
generate_primes
end
@@primes[n]
end
# Doubles the size of the cached prime array and performs the
# Sieve of Eratosthenes on it.
private def generate_primes
new_primes_size = @@primes.size < 1_000_000 ? @@primes.size : 1_000_000
new_primes = Array(Int64).new(new_primes_size) { |i| @@primes.last + (i + 1) * 2 }
new_primes_max = new_primes.last
@@primes.each do |prime|
next if prime == 2
break if prime * prime > new_primes_max
# Here I use the technique of the Sieve of Sundaram. We can
# only test against the odd multiple of the given prime.
# min_composite is the minimum number that is greater than
# the last confirmed prime, and is an odd multiple of
# the given prime.
min_multiple = ((@@primes.last // prime + 1) // 2 * 2 + 1) * prime
min_multiple.step(by: prime * 2, to: new_primes_max) do |multiple|
index = new_primes_size - (new_primes_max - multiple) // 2 - 1
new_primes[index] = 0_i64
end
end
@@primes.concat(new_primes.reject(0_i64))
end
private struct EachDivisor(T)
include Enumerable(T)
def initialize(@exponential_factors : Array(Array(T)))
end
def each(&)
Indexable.each_cartesian(@exponential_factors) do |factors|
yield factors.reduce { |a, b| a * b }
end
end
end
# Returns an enumerator that iterates through the all positive divisors of
# the given number. **The order is not guaranteed.**
# Not in the original Ruby's Prime library.
#
# ```
# AtCoder::Prime.each_divisor(20) do |n|
# puts n
# end # => Puts 1, 2, 4, 5, 10, and 20
#
# AtCoder::Prime.each_divisor(10).map { |n| 1.0 / n }.to_a # => [1.0, 0.5, 0.2, 0.1]
# ```
def each_divisor(value : Int)
raise ArgumentError.new unless value > 0
factors = prime_division(value)
if value == 1
exponential_factors = [[value]]
else
exponential_factors = factors.map do |(factor, count)|
cnt = typeof(value).zero + 1
Array(typeof(value)).new(count + 1) do |i|
cnt_copy = cnt
if i < count
cnt *= factor
end
cnt_copy
end
end
end
EachDivisor(typeof(value)).new(exponential_factors)
end
# :ditto:
def each_divisor(value : T, &block : T ->)
each_divisor(value).each(&block)
end
end
end
struct Int
def prime?
AtCoder::Prime.prime?(self)
end
end
module AtCoder
# Implements [ACL's Math library](https://atcoder.github.io/ac-library/master/document_en/math.html)
module Math
def self.extended_gcd(a, b)
zero = a.class.zero
one = zero + 1
last_remainder, remainder = a.abs, b.abs
x, last_x, y, last_y = zero, one, one, zero
while remainder != 0
quotient, new_remainder = last_remainder.divmod(remainder)
last_remainder, remainder = remainder, new_remainder
x, last_x = last_x - quotient * x, x
y, last_y = last_y - quotient * y, y
end
return last_remainder, last_x * (a < 0 ? -1 : 1)
end
# Implements atcoder::inv_mod(value, modulo).
def self.inv_mod(value, modulo)
gcd, inv = extended_gcd(value, modulo)
if gcd != 1
raise ArgumentError.new("#{value} and #{modulo} are not coprime")
end
inv % modulo
end
# Simplified AtCoder::Math.pow_mod with support of Int64
def self.pow_mod(base, exponent, modulo)
if exponent == 0
return base.class.zero + 1
end
if base == 0
return base
end
b = exponent > 0 ? base : inv_mod(base, modulo)
e = exponent.abs
ret = 1_i64
while e > 0
if e % 2 == 1
ret = mul_mod(ret, b, modulo)
end
b = mul_mod(b, b, modulo)
e //= 2
end
ret
end
# Caluculates a * b % mod without overflow detection
@[AlwaysInline]
def self.mul_mod(a : Int64, b : Int64, mod : Int64)
if mod < Int32::MAX
return a * b % mod
end
# 31-bit width
a_high = (a >> 32).to_u64
# 32-bit width
a_low = (a & 0xFFFFFFFF).to_u64
# 31-bit width
b_high = (b >> 32).to_u64
# 32-bit width
b_low = (b & 0xFFFFFFFF).to_u64
# 31-bit + 32-bit + 1-bit = 64-bit
c = a_high * b_low + b_high * a_low
c_high = c >> 32
c_low = c & 0xFFFFFFFF
# 31-bit + 31-bit
res_high = a_high * b_high + c_high
# 32-bit + 32-bit
res_low = a_low * b_low
res_low_high = res_low >> 32
res_low_low = res_low & 0xFFFFFFFF
# Overflow
if res_low_high + c_low >= 0x100000000
res_high += 1
end
res_low = (((res_low_high + c_low) & 0xFFFFFFFF) << 32) | res_low_low
(((res_high.to_i128 << 64) | res_low) % mod).to_i64
end
@[AlwaysInline]
def self.mul_mod(a, b, mod)
typeof(mod).new(a.to_i64 * b % mod)
end
# Implements atcoder::crt(remainders, modulos).
def self.crt(remainders, modulos)
raise ArgumentError.new unless remainders.size == modulos.size
total_modulo = 1_i64
answer = 0_i64
remainders.zip(modulos).each do |(remainder, modulo)|
gcd, p = extended_gcd(total_modulo, modulo)
if (remainder - answer) % gcd != 0
return 0_i64, 0_i64
end
tmp = (remainder - answer) // gcd * p % (modulo // gcd)
answer += total_modulo * tmp
total_modulo *= modulo // gcd
end
return answer % total_modulo, total_modulo
end
# Implements atcoder::floor_sum(n, m, a, b).
def self.floor_sum(n, m, a, b)
n, m, a, b = n.to_i64, m.to_i64, a.to_i64, b.to_i64
res = 0_i64
if a < 0
a2 = a % m
res -= n * (n - 1) // 2 * ((a2 - a) // m)
a = a2
end
if b < 0
b2 = b % m
res -= n * ((b2 - b) // m)
b = b2
end
res + floor_sum_unsigned(n, m, a, b)
end
private def self.floor_sum_unsigned(n, m, a, b)
res = 0_i64
loop do
if a >= m
res += n * (n - 1) // 2 * (a // m)
a = a % m
end
if b >= m
res += n * (b // m)
b = b % m
end
y_max = a * n + b
break if y_max < m
n = y_max // m
b = y_max % m
m, a = a, m
end
res
end
# Returns `a * b > target`, without concern of overflows.
def self.product_greater_than(a : Int, b : Int, target : Int)
target // b < a
end
def self.get_primitive_root(p : Int)
return 1_i64 if p == 2
n = p - 1
factors = AtCoder::Prime.prime_division(n)
(2_i64..p.to_i64).each do |g|
ok = true
factors.each do |(factor, _)|
if pow_mod(g, n // factor, p) == 1
ok = false
break
end
end
if ok
return g
end
end
raise ArgumentError.new
end
end
end
module AtCoder
# Implements [atcoder::static_modint](https://atcoder.github.io/ac-library/master/document_en/modint.html).
#
# ```
# AtCoder.static_modint(ModInt101, 101_i64)
# alias Mint = AtCoder::ModInt101
# Mint.new(80_i64) + Mint.new(90_i64) # => 89
# ```
macro static_modint(name, modulo)
module AtCoder
# Implements atcoder::modint{{modulo}}.
#
# ```
# alias Mint = AtCoder::{{name}}
# Mint.new(30_i64) // Mint.new(7_i64)
# ```
struct {{name}}
{% if modulo == 998_244_353_i64 %}
MOD = 998_244_353_i64
M = 998_244_353_u32
R = 3_296_722_945_u32
MR = 998_244_351_u32
M2 = 932_051_910_u32
{% elsif modulo == 1_000_000_007_i64 %}
MOD = 1_000_000_007_i64
M = 1_000_000_007_u32
R = 2_068_349_879_u32
MR = 2_226_617_417_u32
M2 = 582_344_008_u32
{% else %}
MOD = {{modulo}}
{% end %}
def self.zero
new
end
@@factorials = Array(self).new
def self.factorial(n)
if @@factorials.empty?
@@factorials = Array(self).new(100_000_i64)
@@factorials << self.new(1)
end
@@factorials.size.upto(n) do |i|
@@factorials << @@factorials.last * i
end
@@factorials[n]
end
def self.permutation(n, k)
raise ArgumentError.new("k cannot be greater than n") unless n >= k
factorial(n) // factorial(n - k)
end
def self.combination(n, k)
raise ArgumentError.new("k cannot be greater than n") unless n >= k
permutation(n, k) // @@factorials[k]
end
def self.repeated_combination(n, k)
combination(n + k - 1, k)
end
def -
self.class.new(0) - self
end
def +
self
end
def +(value)
self + self.class.new(value)
end
def -(value)
self - self.class.new(value)
end
def *(value)
self * self.class.new(value)
end
def /(value)
raise DivisionByZeroError.new if value == 0
self / self.class.new(value)
end
def /(value : self)
raise DivisionByZeroError.new if value == 0
self * value.inv
end
def //(value)
self / value
end
def <<(value)
self * self.class.new(2) ** value
end
def abs
value
end
def pred
self - 1
end
def succ
self + 1
end
def zero?
value == 0
end
def to_i64
value
end
def ==(other : self)
value == other.value
end
def ==(other)
value == other
end
def sqrt
z = self.class.new(1_i64)
until z ** ((MOD - 1) // 2) == MOD - 1
z += 1
end
q = MOD - 1
m = 0
while q.even?
q //= 2
m += 1
end
c = z ** q
t = self ** q
r = self ** ((q + 1) // 2)
m.downto(2) do |i|
tmp = t ** (2 ** (i - 2))
if tmp != 1
r *= c
t *= c ** 2
end
c *= c
end
if r * r == self
r.to_i64 * 2 <= MOD ? r : -r
else
nil
end
end
# ac-library compatibility
def pow(value)
self ** value
end
def val
self.to_i64
end
# ModInt shouldn't be compared
def <(value)
raise NotImplementedError.new("<")
end
def <=(value)
raise NotImplementedError.new("<=")
end
def >(value)
raise NotImplementedError.new(">")
end
def >=(value)
raise NotImplementedError.new(">=")
end
{% if modulo == 998_244_353_i64 || modulo == 1_000_000_007_i64 %}
getter mgy : UInt32
# Initialize using montgomery representation
def self.raw(mgy : UInt32)
ret = new
ret.mgy = mgy
ret
end
def initialize
@mgy = 0
end
def initialize(value : Int)
@mgy = reduce(((value % M).to_u64 + M) * M2)
end
def clone
ret = self.class.new
ret.mgy = @mgy
ret
end
def +(value : self)
ret = self.class.raw(@mgy)
ret.mgy = (ret.mgy.to_i64 + value.mgy - 2*M).to_u32!
if ret.mgy.to_i32! < 0
ret.mgy = (ret.mgy.to_u64 + 2*M).to_u32!
end
ret
end
def -(value : self)
ret = self.class.raw(@mgy)
ret.mgy = (ret.mgy.to_i64 - value.mgy).to_u32!
if ret.mgy.to_i32! < 0
ret.mgy = (ret.mgy.to_u64 + 2*M).to_u32!
end
ret
end
def *(value : self)
ret = self.class.raw(@mgy)
ret.mgy = reduce(ret.mgy.to_u64 * value.mgy)
ret
end
def **(value)
if value == 0
return self.class.new(1)
end
if self.zero?
self
end
b = value > 0 ? self : inv
e = value.abs
ret = self.class.new(1)
while e > 0
if e.odd?
ret *= b
end
b *= b
e >>= 1
end
ret
end
def inv
g, x = AtCoder::Math.extended_gcd(value.to_i32, M.to_i32)
self.class.new(x)
end
def to_s(io : IO)
io << value
end
def inspect(io : IO)
to_s(io)
end
def mgy=(v : UInt32)
@mgy = v
end
@[AlwaysInline]
def reduce(b : UInt64) : UInt32
((b + (b.to_u32!.to_u64 * MR).to_u32!.to_u64 * M) >> 32).to_u32
end
@[AlwaysInline]
def value
ret = reduce(@mgy.to_u64)
ret >= M ? (ret - M).to_i64 : ret.to_i64
end
{% else %}
getter value : Int64
def initialize(@value : Int64 = 0_i64)
@value %= MOD
end
def initialize(value)
@value = value.to_i64 % MOD
end
def clone
self.class.new(@value)
end
def inv
g, x = AtCoder::Math.extended_gcd(@value, MOD)
self.class.new(x)
end
def +(value : self)
self.class.new(@value + value.to_i64)
end
def -(value : self)
self.class.new(@value - value.to_i64)
end
def *(value : self)
self.class.new(@value * value.to_i64)
end
def **(value)
self.class.new(AtCoder::Math.pow_mod(@value, value.to_i64, MOD))
end
delegate to_s, to: @value
delegate inspect, to: @value
{% end %}
end
end
struct Int
def +(value : AtCoder::{{name}})
value + self
end
def -(value : AtCoder::{{name}})
-value + self
end
def *(value : AtCoder::{{name}})
value * self
end
def //(value : AtCoder::{{name}})
value.inv * self
end
def /(value : AtCoder::{{name}})
self // value
end
def ==(value : AtCoder::{{name}})
value == self
end
end
end
end
AtCoder.static_modint(ModInt1000000007, 1_000_000_007_i64)
AtCoder.static_modint(ModInt998244353, 998_244_353_i64)
AtCoder.static_modint(ModInt754974721, 754_974_721_i64)
AtCoder.static_modint(ModInt167772161, 167_772_161_i64)
AtCoder.static_modint(ModInt469762049, 469_762_049_i64)
require "spec"
alias Convolution = AtCoder::Convolution
alias Mint9 = AtCoder::ModInt998244353
describe "Convolution" do
describe ".convolution" do
it "return empty array if any of input array is empty" do
Convolution.convolution([] of Mint9, [] of Mint9).should eq [] of Mint9
Convolution.convolution([Mint9.new(1_i64)], [] of Mint9).should eq [] of Mint9
Convolution.convolution([] of Mint9, [Mint9.new(1_i64)]).should eq [] of Mint9
end
it "calculates convolution of input arrays" do
a = [AtCoder::ModInt998244353.new(1_i64)] * 3
Convolution.convolution(a, a).should eq [1, 2, 3, 2, 1]
end
end
describe ".convolution_ll" do
it "calculates convolution of input arrays" do
a = [1_i64] * 3
Convolution.convolution_ll(a, a).should eq [1, 2, 3, 2, 1]
end
it "can calculate convolutions up to Int64 max" do
a = [1_000_000_000_i64] * 9
Convolution.convolution_ll(a, a).should eq [
1_000_000_000_000_000_000_i64,
2_000_000_000_000_000_000_i64,
3_000_000_000_000_000_000_i64,
4_000_000_000_000_000_000_i64,
5_000_000_000_000_000_000_i64,
6_000_000_000_000_000_000_i64,
7_000_000_000_000_000_000_i64,
8_000_000_000_000_000_000_i64,
9_000_000_000_000_000_000_i64,
8_000_000_000_000_000_000_i64,
7_000_000_000_000_000_000_i64,
6_000_000_000_000_000_000_i64,
5_000_000_000_000_000_000_i64,
4_000_000_000_000_000_000_i64,
3_000_000_000_000_000_000_i64,
2_000_000_000_000_000_000_i64,
1_000_000_000_000_000_000_i64,
]
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/dsu.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements [atcoder::dsu](https://atcoder.github.io/ac-library/master/document_en/dsu.html).
#
# ```
# dsu = AtCoder::DSU.new(10)
# dsu.merge(0, 2)
# dsu.merge(4, 2)
# dsu.same?(0, 4) # => true
# dsu.size(4) # => 3
# ```
class DSU
getter parents : Array(Int64)
getter sizes : Array(Int64)
getter size : Int64
def initialize(@size)
@parents = Array.new(size, &.to_i64)
@sizes = Array.new(size, 1_i64)
end
# Implements atcoder::dsu.leader(node).
def leader(node)
until @parents[node] == node
@parents[node] = @parents[@parents[node]]
node = @parents[node]
end
node
end
# Implements atcoder::dsu.merge(a, b).
def merge(a, b)
leader_a = leader(a.to_i64)
leader_b = leader(b.to_i64)
unless leader_a == leader_b
if @sizes[leader_a] < @sizes[leader_b]
leader_a, leader_b = leader_b, leader_a
end
@parents[leader_b] = leader_a
@sizes[leader_a] += @sizes[leader_b]
end
end
# Implements atcoder::dsu.same(a, b).
def same?(a, b)
leader(a) == leader(b)
end
# Implements atcoder::dsu.size().
def size(node)
@sizes[leader(node)]
end
# Implements atcoder::dsu.groups().
# This method returns `Set` instead of list.
def groups
groups = Hash(Int64, Set(Int64)).new { |h, k| h[k] = Set(Int64).new }
@size.times do |i|
groups[leader(i)] << i
end
groups.values.to_set
end
end
end
require "spec"
alias DSU = AtCoder::DSU
describe "DSU" do
describe "#merge" do
it "unifies group of trees into unions" do
tree = DSU.new(9_i64)
tree.merge(0, 2)
tree.merge(2, 4)
tree.merge(4, 6)
tree.merge(7, 1)
tree.merge(7, 3)
tree.merge(7, 5)
tree.same?(2, 6).should eq true
tree.same?(4, 2).should eq true
tree.same?(1, 3).should eq true
tree.same?(3, 5).should eq true
tree.same?(8, 8).should eq true
tree.same?(2, 1).should eq false
tree.same?(7, 0).should eq false
tree.same?(5, 4).should eq false
tree.same?(8, 7).should eq false
tree.size(4).should eq 4
tree.size(1).should eq 4
tree.size(8).should eq 1
tree.groups.should eq Set{
Set{0, 2, 4, 6},
Set{1, 3, 5, 7},
Set{8},
}
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/fenwick_tree.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements [atcoder::fenwick_tree](https://atcoder.github.io/ac-library/master/document_en/fenwicktree.html).
#
# ```
# tree = AtCoder::FenwickTree(Int64).new(10)
# tree.add(3, 10)
# tree.add(5, 20)
# tree[3..5] # => 30
# tree[3...5] # => 10
# ```
class FenwickTree(T)
getter size : Int64
getter bits : Array(T)
def initialize(@size : Int64)
@bits = Array(T).new(@size, T.zero)
end
def initialize(@bits : Array)
@bits = @bits.dup
@size = @bits.size.to_i64
(1...@size).each do |index|
up = index + (index & -index)
next if up > @size
@bits[up - 1] += @bits[index - 1]
end
end
# Implements atcoder::fenwick_tree.add(index, value)
def add(index, value)
index += 1 # convert to 1-indexed
while index <= @size
@bits[index - 1] += value
index += index & -index
end
end
# Exclusive left sum
def left_sum(index)
ret = T.zero
while index >= 1
ret += @bits[index - 1]
index -= index & -index
end
ret
end
# Implements atcoder::fenwick_tree.sum(left, right)
def sum(left, right)
left_sum(right) - left_sum(left)
end
# Implements atcoder::fenwick_tree.sum(left, right)
#
# Open ended ranges are clamped at the start and end of array, respectively.
def sum(range : Range)
left = range.begin || 0
right = range.exclusive? ? (range.end || @size) : (range.end || @size - 1) + 1
sum(left, right)
end
# :ditto:
def [](range : Range)
sum(range)
end
end
end
# require "../src/mod_int.cr"
# require "./spec_helper.cr"
# require "../src/mod_int.cr"
AtCoder.static_modint(ModInt7, 7_i64)
alias Mint = AtCoder::ModInt1000000007
require "spec"
alias FenwickTree = AtCoder::FenwickTree
describe "FenwickTree" do
describe "#new(array)" do
it "equals #new(int64)" do
a = [10, 20, 30, 40, 50, 60, 70].map(&.to_i64)
tree1 = FenwickTree(Int64).new(a)
tree2 = FenwickTree(Int64).new(a.size.to_i64)
a.each_with_index { |e, i| tree2.add(i, e) }
tree1.size.should eq tree2.size
tree1.bits.should eq tree2.bits
end
end
describe "#left_sum" do
it "calculates left_sum of given list" do
tree = FenwickTree(Int64).new(11_i64)
11_i64.times do |i|
tree.add(i, i)
end
tree.left_sum(0).should eq 0
tree.left_sum(1).should eq 0
tree.left_sum(2).should eq 1
tree.left_sum(3).should eq 3
tree.left_sum(4).should eq 6
tree.left_sum(5).should eq 10
tree.left_sum(6).should eq 15
tree.left_sum(7).should eq 21
tree.left_sum(8).should eq 28
tree.left_sum(9).should eq 36
tree.left_sum(10).should eq 45
tree.left_sum(11).should eq 55
end
end
describe "#[]" do
it "calculates range sum of given list" do
tree = FenwickTree(Int64).new(11_i64)
11_i64.times do |i|
tree.add(i, i)
end
tree[2..4].should eq 9
tree[0..10].should eq 55
tree[0...3].should eq 3
tree[3...7].should eq 18
tree[0...11].should eq 55
tree[7...7].should eq 0
tree[3...].should eq 52
tree[...6].should eq 15
tree[...].should eq 55
tree[3..].should eq 52
tree[..6].should eq 21
tree[..].should eq 55
end
end
it "can be used with ModInt" do
tree = FenwickTree(ModInt7).new(11_i64)
11_i64.times do |i|
tree.add(i, i)
end
tree[4..4].should eq 4
tree[10..10].should eq 3
tree[4...7].should eq 1
tree[0...11].should eq 6
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/lazy_seg_tree.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements [atcoder::lazy_segtree](https://atcoder.github.io/ac-library/master/document_en/lazysegtree.html).
#
# The identity element will be implicitly defined as nil, so you don't
# have to manually define it. In the other words, you cannot include
# nil into an element of the monoid.
#
# Similarly, the identity map of F will be implicitly defined as nil,
# so you don't have to manually define it. In the other words, you
# cannot include nil into an element of the set F.
#
# ```
# op = ->(a : Int32, b : Int32) { [a, b].min }
# mapping = ->(f : Int32, x : Int32) { f }
# composition = ->(a : Int32, b : Int32) { a }
# tree = AtCoder::LazySegTree(Int32, Int32).new((0...100).to_a, op, mapping, composition)
# tree[10...50] # => 10
# tree[20...60] = 0
# tree[50...80] # => 0
# ```
class LazySegTree(S, F)
@n : Int32
@height : Int32
@n_leaves : Int32
def initialize(values : Array(S), @operator : S, S -> S, @application : F, S -> S, @composition : F, F -> F)
@n = values.size
@height = log2_ceil(@n)
@n_leaves = 1 << @height
@segments = Array(S | Nil).new(2 * @n_leaves, nil)
@applicators = Array(F | Nil).new(@n_leaves, nil)
# initialize segments
values.each_with_index { |x, i| @segments[@n_leaves + i] = x.as(S | Nil) }
(@n_leaves - 1).downto(1) { |i| refresh(i) }
end
# Implements atcoder::lazy_segtree.set(index, applicator).
def set(index : Int, value : S)
index += @n_leaves
@height.downto(1) { |j| propagate(ancestor(index, j)) }
@segments[index] = value.as(S | Nil)
(1..@height).each { |j| refresh(ancestor(index, j)) }
end
# Implements atcoder::lazy_segtree.apply(index, applicator).
def []=(index : Int, applicator : F)
index += @n_leaves
@height.downto(1) { |j| propagate(ancestor(index, j)) }
@segments[index] = apply(applicator, @segments[index])
(1..@height).each { |j| refresh(ancestor(index, j)) }
applicator
end
# Implements atcoder::lazy_segtree.apply(left, right, applicator).
# ameba:disable Metrics/CyclomaticComplexity
def []=(range : Range, applicator : F)
l = (range.begin || 0) + @n_leaves
r = (range.exclusive? ? (range.end || @n_leaves) : (range.end || @n_leaves - 1) + 1) + @n_leaves
@height.downto(1) do |i|
propagate(ancestor(l, i)) if right_side_child?(l, i)
propagate(ancestor(r - 1, i)) if right_side_child?(r, i)
end
l2, r2 = l, r
while l2 < r2
if l2.odd?
all_apply(l2, applicator)
l2 += 1
end
if r2.odd?
r2 -= 1
all_apply(r2, applicator)
end
l2 >>= 1
r2 >>= 1
end
(1..@height).each do |i|
refresh(ancestor(l, i)) if right_side_child?(l, i)
refresh(ancestor(r - 1, i)) if right_side_child?(r, i)
end
applicator
end
# Implements atcoder::lazy_segtree.get(index).
def [](index : Int)
index += @n_leaves
@height.downto(1) { |j| propagate(ancestor(index, j)) }
@segments[index].not_nil!
end
# Implements atcoder::lazy_segtree.prod(left, right).
def [](range : Range)
l = (range.begin || 0) + @n_leaves
r = (range.exclusive? ? (range.end || @n_leaves) : (range.end || @n_leaves - 1) + 1) + @n_leaves
@height.downto(1) do |i|
propagate(ancestor(l, i)) if right_side_child?(l, i)
propagate(ancestor(r - 1, i)) if right_side_child?(r, i)
end
sml, smr = nil.as(S | Nil), nil.as(S | Nil)
while l < r
if l.odd?
sml = operate(sml, @segments[l])
l += 1
end
if r.odd?
r -= 1
smr = operate(@segments[r], smr)
end
l >>= 1
r >>= 1
end
operate(sml, smr).not_nil!
end
# Implements atcoder::lazy_segtree.all_prod().
def all_prod
self[0...@n]
end
# Implements atcoder::lazy_segtree.max_right(left, g).
def max_right(left, e : S | Nil = nil, & : S -> Bool)
unless 0 <= left && left <= @n
raise IndexError.new("{left: #{left}} must greater than or equal to 0 and less than or equal to {n: #{@n}}")
end
unless e.nil?
return nil unless yield e
end
return @n if left == @n
left += @n_leaves
@height.downto(1) { |i| propagate(ancestor(left, i)) }
sm = e
loop do
while left.even?
left >>= 1
end
res = operate(sm, @segments[left])
unless res.nil? || yield res
while left < @n_leaves
propagate(left)
left = 2*left
res = operate(sm, @segments[left])
if res.nil? || yield res
sm = res
left += 1
end
end
return left - @n_leaves
end
sm = operate(sm, @segments[left])
left += 1
ffs = left & -left
break if ffs == left
end
@n
end
# Implements atcoder::lazy_segtree.min_left(right, g).
def min_left(right, e : S | Nil = nil, & : S -> Bool)
unless 0 <= right && right <= @n
raise IndexError.new("{right: #{right}} must greater than or equal to 0 and less than or equal to {n: #{@n}}")
end
unless e.nil?
return nil unless yield e
end
return 0 if right == 0
right += @n_leaves
@height.downto(1) { |i| propagate(ancestor(right - 1, i)) }
sm = e
loop do
right -= 1
while right > 1 && right.odd?
right >>= 1
end
res = operate(@segments[right], sm)
unless res.nil? || yield res
while right < @n_leaves
propagate(right)
right = 2*right + 1
res = operate(@segments[right], sm)
if res.nil? || yield res
sm = res
right -= 1
end
end
return right + 1 - @n_leaves
end
sm = operate(@segments[right], sm)
ffs = right & -right
break if ffs == right
end
0
end
@[AlwaysInline]
private def operate(a : S | Nil, b : S | Nil)
if a.nil?
b
elsif b.nil?
a
else
@operator.call(a, b)
end
end
@[AlwaysInline]
private def apply(f : F | Nil, x : S | Nil)
if f.nil?
x
elsif x.nil?
nil
else
@application.call(f, x)
end
end
@[AlwaysInline]
private def compose(a : F | Nil, b : F | Nil)
if a.nil?
b
elsif b.nil?
a
else
@composition.call(a, b)
end
end
@[AlwaysInline]
private def refresh(node : Int)
child1 = 2*node
child2 = 2*node + 1
@segments[node] = operate(@segments[child1], @segments[child2])
end
@[AlwaysInline]
private def all_apply(node, applicator : F | Nil)
@segments[node] = apply(applicator, @segments[node])
unless leaf?(node)
@applicators[node] = compose(applicator, @applicators[node])
end
end
@[AlwaysInline]
private def propagate(node : Int)
child1 = 2*node
child2 = 2*node + 1
all_apply(child1, @applicators[node])
all_apply(child2, @applicators[node])
@applicators[node] = nil
end
@[AlwaysInline]
private def right_side_child?(child, n_gens_ago)
((child >> n_gens_ago) << n_gens_ago) != child
end
@[AlwaysInline]
private def ancestor(node, n_gens_ago)
node >> n_gens_ago
end
@[AlwaysInline]
private def leaf?(node)
node >= @n_leaves
end
@[AlwaysInline]
private def log2_ceil(n : Int32) : Int32
sizeof(Int32)*8 - (n - 1).leading_zeros_count
end
end
end
# require "../src/mod_int.cr"
# require "./spec_helper.cr"
require "spec"
alias LazySegTree = AtCoder::LazySegTree
describe "LazySegTree" do
describe "#[]" do
it "simply returns value of tree" do
op = ->(a : Int32, b : Int32) { a + b }
mapping = ->(f : Int32, x : Int32) { x * f }
composition = ->(a : Int32, b : Int32) { a * b }
segtree = LazySegTree(Int32, Int32).new((1..100).to_a, op, mapping, composition)
segtree[0].should eq 1
segtree[10].should eq 11
segtree[19].should eq 20
segtree[50].should eq 51
segtree[99].should eq 100
end
it "takes sum of selected range" do
op = ->(a : Int32, b : Int32) { a + b }
mapping = ->(f : Int32, x : Int32) { x * f }
composition = ->(a : Int32, b : Int32) { a * b }
segtree = LazySegTree(Int32, Int32).new((1..100).to_a, op, mapping, composition)
segtree[0...100].should eq 5050
segtree[10...50].should eq 1220
segtree[50...70].should eq 1210
segtree[12...33].should eq 483
segtree[0..0].should eq 1
segtree[16..16].should eq 17
segtree[49..49].should eq 50
segtree[88..88].should eq 89
segtree[10...].should eq 4995
segtree[...60].should eq 1830
segtree[...].should eq 5050
segtree[10..].should eq 4995
segtree[..60].should eq 1891
segtree[..].should eq 5050
end
end
describe "#set" do
it "sets value at index" do
op = ->(a : Int32, b : Int32) { a + b }
mapping = ->(f : Int32, x : Int32) { x * f }
composition = ->(a : Int32, b : Int32) { a * b }
segtree = LazySegTree(Int32, Int32).new([0] * 100, op, mapping, composition)
100.times { |i| segtree.set(i, i + 1) }
segtree[0...100].should eq 5050
segtree[10...50].should eq 1220
segtree[50...70].should eq 1210
segtree[12...33].should eq 483
segtree.set(0, 1000000)
segtree.set(10, 1000000)
segtree.set(100, 1000000)
segtree[0...100].should eq 5050 + 2 * 1000000 - (1 + 11)
segtree[0...101].should eq 5151 + 3 * 1000000 - (1 + 11 + 101)
end
end
describe "#[]=" do
it "updates value at selected index" do
op = ->(a : Int32, b : Int32) { a + b }
mapping = ->(f : Int32, x : Int32) { x * f }
composition = ->(a : Int32, b : Int32) { a * b }
segtree = LazySegTree(Int32, Int32).new((1..100).to_a, op, mapping, composition)
segtree[40] = 9999
segtree[40].should eq 41 * 9999
segtree[0...100].should eq 5050 + 41 * 9998
segtree[99] = 99999
segtree[99].should eq 9999900
segtree[0...100].should eq 5050 + 41 * 9998 - 100 + 9999900
end
it "bulk-updates values" do
op = ->(a : Int32, b : Int32) { a + b }
mapping = ->(f : Int32, x : Int32) { x * f }
composition = ->(a : Int32, b : Int32) { a * b }
segtree = LazySegTree(Int32, Int32).new((1..100).to_a, op, mapping, composition)
segtree[10...30] = 2
segtree[9].should eq 10
segtree[10].should eq 22
segtree[19].should eq 40
segtree[29].should eq 60
segtree[30].should eq 31
segtree[49].should eq 50
segtree[99].should eq 100
segtree[0...100].should eq 5050 + 410
segtree[29...40] = 5
segtree[28].should eq 58
segtree[29].should eq 300
segtree[30].should eq 155
segtree[39].should eq 200
segtree[40].should eq 41
segtree[41].should eq 42
segtree[0...100].should eq 5050 + 410 + 240 + 355 * 4
segtree[1...] = 2
segtree[...5] = 2
segtree[...] = 2
segtree[0].should eq 1 * 2**2
segtree[1].should eq 2 * 2**3
segtree[2].should eq 3 * 2**3
segtree[3].should eq 4 * 2**3
segtree[4].should eq 5 * 2**3
segtree[5].should eq 6 * 2**2
segtree[6].should eq 7 * 2**2
segtree[99].should eq 100 * 2**2
segtree[1..] = 2
segtree[..5] = 2
segtree[..] = 2
segtree[0].should eq 1 * 2**4
segtree[1].should eq 2 * 2**6
segtree[2].should eq 3 * 2**6
segtree[3].should eq 4 * 2**6
segtree[4].should eq 5 * 2**6
segtree[5].should eq 6 * 2**5
segtree[6].should eq 7 * 2**4
segtree[99].should eq 100 * 2**4
end
end
describe "#max_right" do
it "applies a binary search on the segment tree" do
op = ->(a : Int32, b : Int32) { Math.min(a, b) }
mapping = ->(f : Int32, x : Int32) { x + f }
composition = ->(a : Int32, b : Int32) { a + b }
segtree = LazySegTree(Int32, Int32).new((-10..10).map(&.**(2)), op, mapping, composition)
(0..20).map { |i| segtree[0..i] }.to_a.should eq [100, 81, 64, 49, 36, 25, 16, 9, 4, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
segtree.max_right(0) { |x| x >= 10**9 }.should eq 0
segtree.max_right(0) { |x| x >= 20 }.should eq 6
segtree.max_right(0) { |x| x >= 16 }.should eq 7
segtree.max_right(0) { |x| x >= -1 }.should eq 21
segtree.max_right(0) { |x| x <= 10**9 }.should eq 21
segtree.max_right(0, e: Int32::MAX) { |x| x >= 10**9 }.should eq 0
segtree.max_right(0, e: Int32::MAX) { |x| x <= 10**9 }.should eq nil
end
end
describe "#min_left" do
it "applies a binary search on the segment tree" do
op = ->(a : Int32, b : Int32) { Math.min(a, b) }
mapping = ->(f : Int32, x : Int32) { x + f }
composition = ->(a : Int32, b : Int32) { a + b }
segtree = LazySegTree(Int32, Int32).new((-10..10).map(&.**(2)), op, mapping, composition)
(0..20).map { |i| segtree[i..20] }.to_a.should eq [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
segtree.min_left(21) { |x| x >= 10**9 }.should eq 21
segtree.min_left(21) { |x| x >= 20 }.should eq 15
segtree.min_left(21) { |x| x >= 16 }.should eq 14
segtree.min_left(21) { |x| x >= -1 }.should eq 0
segtree.min_left(21) { |x| x <= 10**9 }.should eq 0
segtree.min_left(21, e: Int32::MAX) { |x| x <= 10**9 }.should eq nil
segtree.min_left(21, e: Int32::MAX) { |x| x >= 10**9 }.should eq 21
end
end
it "can be used with ModInt" do
op = ->(a : Mint, b : Mint) { a + b }
mapping = ->(f : Int64, x : Mint) { x * f }
composition = ->(a : Int64, b : Int64) { a * b }
values = 100.times.map { |i| Mint.new(1_i64) << i }.to_a
segtree = LazySegTree(Mint, Int64).new(values, op, mapping, composition)
segtree[0...32].should eq 0xFFFFFFFF % Mint::MOD
segtree[0...100] = 8
segtree[0...100] = 32
segtree[8...40].should eq 0xFFFFFFFF0000 % Mint::MOD
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/math.cr"
# require "./spec_helper.cr"
require "spec"
def floor_sum_naive(n, m, a, b)
sum = 0i64
n.times do |i|
z = a * i + b
sum += (z - z % m) // m
end
sum
end
describe "Math" do
describe ".inv_mod" do
it "generates inverse number under modulo" do
AtCoder::Math.inv_mod(1, 9).should eq 1
AtCoder::Math.inv_mod(8, 13).should eq 5
AtCoder::Math.inv_mod(3, 299).should eq 100
AtCoder::Math.inv_mod(-1, 7).should eq 6
AtCoder::Math.inv_mod(1_000_000_007_i64, 1_000_000_000_000_000_000_i64).should eq 551_020_408_142_857_143_i64
end
it "raises error when there's no inverse" do
expect_raises(ArgumentError) { AtCoder::Math.inv_mod(3, 9) }
expect_raises(ArgumentError) { AtCoder::Math.inv_mod(0, 1000) }
end
end
describe ".pow_mod" do
it "generates powered number" do
AtCoder::Math.pow_mod(10, 3, 29).should eq 14
AtCoder::Math.pow_mod(123_456, 0, 1_000_003).should eq 1
AtCoder::Math.pow_mod(100, 100, 299).should eq 100
AtCoder::Math.pow_mod(3, -100, 299).should eq 100
AtCoder::Math.pow_mod(0, 0, 1000).should eq 1
AtCoder::Math.pow_mod(100, 10, 1000).should eq 0
AtCoder::Math.pow_mod(10, 2, 1000).should eq 100
AtCoder::Math.pow_mod(100, 0, 1000).should eq 1
AtCoder::Math.pow_mod(0, -10, 1000).should eq 0
AtCoder::Math.pow_mod(123_456_789_i64, 1_000_000_006_i64, 1_000_000_007_i64).should eq 1
end
it "raises error when there's no inverse" do
expect_raises(ArgumentError) { AtCoder::Math.pow_mod(3, -2, 9) }
expect_raises(ArgumentError) { AtCoder::Math.pow_mod(100, -10, 1000) }
end
end
describe ".crt" do
it "calculates chinese remainder theorem answers" do
AtCoder::Math.crt([5, 7], [10, 11]).should eq({95, 110})
AtCoder::Math.crt([73, 10073], [10000, 20000]).should eq({10073, 20000})
AtCoder::Math.crt([2, 3, 2], [3, 5, 7]).should eq({23, 105})
AtCoder::Math.crt([10, 15, 10], [15, 25, 35]).should eq({115, 525})
AtCoder::Math.crt([11, 16, 11], [15, 25, 35]).should eq({116, 525})
end
it "returns {0, 1} if the given array is empty" do
AtCoder::Math.crt([] of Int32, [] of Int32).should eq({0, 1})
end
it "returns {0, 0} if there's no such answer" do
AtCoder::Math.crt([1, 2], [10, 15]).should eq({0, 0})
AtCoder::Math.crt([109, 185], [213, 267]).should eq({0, 0})
end
it "raises error when sizes of given arrays don't match" do
expect_raises(ArgumentError) { AtCoder::Math.crt([] of Int32, [1]) }
expect_raises(ArgumentError) { AtCoder::Math.crt([1, 2, 3], [4, 5]) }
end
end
describe ".floor_sum" do
it "calculates floor_sum" do
# https://atcoder.jp/contests/practice2/tasks/practice2_c
AtCoder::Math.floor_sum(4, 10, 6, 3).should eq 3
AtCoder::Math.floor_sum(6, 5, 4, 3).should eq 13
AtCoder::Math.floor_sum(1, 1, 0, 0).should eq 0
AtCoder::Math.floor_sum(31415, 92653, 58979, 32384).should eq 314_095_480
AtCoder::Math.floor_sum(1_000_000_000, 1_000_000_000, 999_999_999, 999_999_999).should eq 499_999_999_500_000_000
end
it "equals floor_sum_naive" do
(0...20).each do |n|
(1...20).each do |m|
(-20...20).each do |a|
(-20...20).each do |b|
AtCoder::Math.floor_sum(n, m, a, b).should eq floor_sum_naive(n, m, a, b)
end
end
end
end
end
end
describe ".product_greater_than" do
it "checks if the product of a and b is greater than target" do
AtCoder::Math.product_greater_than(4, 3, 13).should eq false
AtCoder::Math.product_greater_than(4, 3, 12).should eq false
AtCoder::Math.product_greater_than(4, 3, 11).should eq true
AtCoder::Math.product_greater_than(4, 3, 10).should eq true
AtCoder::Math.product_greater_than(3, 4, 13).should eq false
AtCoder::Math.product_greater_than(3, 4, 12).should eq false
AtCoder::Math.product_greater_than(3, 4, 11).should eq true
AtCoder::Math.product_greater_than(3, 4, 10).should eq true
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/min_cost_flow.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./graph.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "./priority_queue.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements standard priority queue like [std::priority_queue](https://en.cppreference.com/w/cpp/container/priority_queue).
#
# ```
# q = AtCoder::PriorityQueue(Int64).new
# q << 1_i64
# q << 3_i64
# q << 2_i64
# q.pop # => 3
# q.pop # => 2
# q.pop # => 1
# ```
class PriorityQueue(T)
include Enumerable(T)
getter heap : Array(T)
# Create a new queue in ascending order of priority.
def self.max
self.new { |a, b| a <= b }
end
# Create a new queue in ascending order of priority with the elements in *enumerable*.
def self.max(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a <= b }
end
# Create a new queue in descending order of priority.
def self.min
self.new { |a, b| a >= b }
end
# Create a new queue in descending order of priority with the elements in *enumerable*.
def self.min(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a >= b }
end
def initialize
initialize { |a, b| a <= b }
end
# Initializes queue with the elements in *enumerable*.
def self.new(enumerable : Enumerable(T))
self.new(enumerable) { |a, b| a <= b }
end
# Initializes queue with the custom comperator.
#
# If the second argument `b` should be popped earlier than
# the first argument `a`, return `true`. Else, return `false`.
#
# ```
# q = AtCoder::PriorityQueue(Int64).new { |a, b| a >= b }
# q << 1_i64
# q << 3_i64
# q << 2_i64
# q.pop # => 1
# q.pop # => 2
# q.pop # => 3
# ```
def initialize(&block : T, T -> Bool)
@heap = Array(T).new
@compare_proc = block
end
# Initializes queue with the elements in *enumerable* and the custom comperator.
#
# If the second argument `b` should be popped earlier than
# the first argument `a`, return `true`. Else, return `false`.
#
# ```
# q = AtCoder::PriorityQueue.new([1, 3, 2]) { |a, b| a >= b }
# q.pop # => 1
# q.pop # => 2
# q.pop # => 3
# ```
def initialize(enumerable : Enumerable(T), &block : T, T -> Bool)
@heap = enumerable.to_a
@compare_proc = block
len = @heap.size
(len // 2 - 1).downto(0) do |parent|
v = @heap[parent]
child = parent * 2 + 1
while child < len
if child + 1 < len && @compare_proc.call(@heap[child], @heap[child + 1])
child += 1
end
unless @compare_proc.call(v, @heap[child])
break
end
@heap[parent] = @heap[child]
parent = child
child = parent * 2 + 1
end
@heap[parent] = v
end
end
# Pushes value into the queue.
# This method returns self, so several calls can be chained.
def push(v : T) : self
@heap << v
index = @heap.size - 1
while index != 0
parent = (index - 1) // 2
if @compare_proc.call(@heap[index], @heap[parent])
break
end
@heap[parent], @heap[index] = @heap[index], @heap[parent]
index = parent
end
self
end
# Alias of `push`
def <<(v : T) : self
push(v)
end
# Pops value from the queue.
def pop
if @heap.size == 0
return nil
end
if @heap.size == 1
return @heap.pop
end
ret = @heap.first
@heap[0] = @heap.pop
index = 0
while index * 2 + 1 < @heap.size
child = if index * 2 + 2 < @heap.size && !@compare_proc.call(@heap[index * 2 + 2], @heap[index * 2 + 1])
index * 2 + 2
else
index * 2 + 1
end
if @compare_proc.call(@heap[child], @heap[index])
break
end
@heap[child], @heap[index] = @heap[index], @heap[child]
index = child
end
ret
end
# Pops value from the queue.
# Raises `Enumerable::EmptyError` if queue is of 0 size.
def pop!
pop || raise Enumerable::EmptyError.new
end
# Yields each item in the queue in comparator's order.
def each(&)
@heap.sort { |a, b| @compare_proc.call(a, b) ? 1 : -1 }.each do |e|
yield e
end
end
# Returns, but does not remove, the head of the queue.
def first(&)
@heap.first { yield }
end
# Returns `true` if the queue is empty.
delegate :empty?, to: @heap
# Returns size of the queue.
delegate :size, to: @heap
end
end
module AtCoder
class Graph(NodeInfo, EdgeInfo)
@size_bits : Int32
getter visited : Set(Int64)
def initialize(@nodes : Array(NodeInfo))
@size = @nodes.size.to_i64
@edges = [] of EdgeInfo
@adjacencies = Array(Hash(Int64, Int64)).new(@size) { Hash(Int64, Int64).new }
@visited = Set(Int64).new
@size_bits = @size.to_s(2).size
end
def initialize(@size : Int64, initial_node : NodeInfo = nil)
@nodes = Array(NodeInfo).new(@size, initial_node)
@edges = [] of EdgeInfo
@adjacencies = Array(Hash(Int64, Int64)).new(@size) { Hash(Int64, Int64).new }
@visited = Set(Int64).new
@size_bits = @size.to_s(2).size
end
# Performs Dijkstra's Algorithm to calculate the distance of each node from `start_node`.
# To use this method, `EdgeInfo` must implement `.zero` and `#+(EdgeInfo)` and `#>(EdgeInfo)`.
def dijkstra(start_node)
dist = Array(EdgeInfo | Nil).new(@size, nil)
dist[start_node] = EdgeInfo.zero
prev_nodes = Array(Int64 | Nil).new(@size, nil)
queue = AtCoder::PriorityQueue({EdgeInfo, Int64}).new { |(ad, av), (bd, bv)| ad > bd }
queue << {EdgeInfo.zero, start_node.to_i64}
until queue.empty?
d, v = queue.pop.not_nil!
current_dist = dist[v].not_nil!
next if d > current_dist
@adjacencies[v].each do |(to, edge_id)|
cost = @edges[edge_id]
target_dist = dist[to]
if target_dist.nil? || target_dist > current_dist + cost
dist[to] = current_dist + cost
prev_nodes[to] = v
queue << {current_dist + cost, to}
end
end
end
Array({dist: EdgeInfo | Nil, prev: Int64 | Nil}).new(@size) do |i|
{dist: dist[i], prev: prev_nodes[i]}
end
end
def dfs(node : Int64, initial_value : T, &block : Int64, T, NamedTuple(
node: Int64,
node_info: NodeInfo | Nil,
edge: Int64 | Nil,
edge_info: EdgeInfo | Nil,
parent: Int64,
), (T ->) ->) forall T
@visited = Set(Int64).new
info = {
node: node,
node_info: @nodes[node].as(NodeInfo | Nil),
edge: nil.as(Int64 | Nil),
edge_info: nil.as(EdgeInfo | Nil),
parent: -1_i64,
}
block.call(node, initial_value, info, ->(new_value : T) {
dfs(node, -1_i64, new_value, &block)
})
end
private def dfs(node : Int64, parent : Int64, value : T, &block : Int64, T, NamedTuple(
node: Int64,
node_info: NodeInfo | Nil,
edge: Int64 | Nil,
edge_info: EdgeInfo | Nil,
parent: Int64,
), (T ->) ->) forall T
@visited << node
@adjacencies[node].each do |(child, edge)|
next if @visited.includes?(child)
info = {
node: child,
node_info: @nodes[child].as(NodeInfo | Nil),
edge: edge.as(Int64 | Nil),
edge_info: @edges[edge].as(EdgeInfo | Nil),
parent: node,
}
block.call(child, value, info, ->(new_value : T) {
dfs(child, node, new_value, &block)
})
end
end
end
class DirectedGraph(NodeInfo, EdgeInfo) < Graph(NodeInfo, EdgeInfo)
def add_edge(from, to, edge : EdgeInfo = 1_i64)
@edges << edge
edge_id = @edges.size.to_i64 - 1
@adjacencies[from][to.to_i64] = edge_id
edge_id
end
def get_edge(from, to)
edge_id = @adjacencies[from][to.to_i64]
@edges[edge_id]
end
def update_edge(from, to, edge : EdgeInfo = 1_i64)
edge_id = @adjacencies[from][to.to_i64]
@edges[edge_id] = edge
edge_id
end
end
class UndirectedGraph(NodeInfo, EdgeInfo) < Graph(NodeInfo, EdgeInfo)
def add_edge(a, b, edge : EdgeInfo = 1_i64)
@edges << edge
edge_id = @edges.size.to_i64 - 1
@adjacencies[a][b.to_i64] = edge_id
@adjacencies[b][a.to_i64] = edge_id
edge_id
end
def get_edge(a, b)
edge_id = @adjacencies[a][b.to_i64]
@edges[edge_id]
end
def update_edge(a, b, edge : EdgeInfo = 1_i64)
edge_id = @adjacencies[a][b.to_i64]
@edges[edge_id] = edge
edge_id
end
end
class Tree(NodeInfo, EdgeInfo) < UndirectedGraph(NodeInfo, EdgeInfo)
@lca_root : Int64 | Nil
@lca_ancestors : Array(Array(Int64)) | Nil
@lca_depths : Array(Int64) | Nil
def diameter
@farthest_node = -1_i64
@farthest_depth = 0_i64
dfs(0_i64, 0_i64) do |node, depth, info, callback|
weight = info[:edge_info]
depth += weight.nil? ? 0 : weight
if @farthest_depth.not_nil! < depth
@farthest_node = node
@farthest_depth = depth
end
callback.call(depth)
end
start_node = @farthest_node.not_nil!
@farthest_node = -1_i64
@farthest_depth = 0_i64
@parents = Array(Int64).new(@size, -1_i64)
dfs(start_node, 0_i64) do |node, depth, info, callback|
weight = info[:edge_info]
depth += weight.nil? ? 0 : weight
@parents.not_nil![node] = info[:parent]
if @farthest_depth.not_nil! < depth
@farthest_node = node
@farthest_depth = depth
end
callback.call(depth)
end
{@farthest_depth.not_nil!, start_node, @farthest_node.not_nil!, @parents.not_nil!}
end
private def lca_precompute(root)
lca_ancestors = Array(Array(Int64)).new(@size) { Array(Int64).new(@size_bits, -1_i64) }
lca_depths = Array(Int64).new(@size, -1_i64)
dfs(root, 0_i64) do |node, depth, info, callback|
lca_ancestors[node][0] = info[:parent]
lca_depths[node] = depth
callback.call(depth + 1)
end
1.upto(@size_bits - 1) do |i|
@size.times do |node|
if lca_ancestors[node][i - 1] == -1
lca_ancestors[node][i] = -1_i64
else
lca_ancestors[node][i] = lca_ancestors[lca_ancestors[node][i - 1]][i - 1]
end
end
end
@lca_root = root
@lca_ancestors = lca_ancestors
@lca_depths = lca_depths
end
private def lca_nth_prev(node, dist)
lca_ancestors = @lca_ancestors.not_nil!
i = 0_i64
until dist == 0
if dist.odd?
node = lca_ancestors[node][i]
if node == -1
raise Exception.new("#{dist}th previous node of #{node} does not exist")
end
end
dist //= 2
i += 1
end
node
end
# ameba:disable Metrics/CyclomaticComplexity
def lca(a, b)
if @lca_root.nil?
lca_precompute(0_i64)
end
lca_ancestors = @lca_ancestors.not_nil!
lca_depths = @lca_depths.not_nil!
depth_a = lca_depths[a]
depth_b = lca_depths[b]
if depth_a < depth_b
depth_a, depth_b, a, b = depth_b, depth_a, b, a
end
if depth_a != depth_b
a = lca_nth_prev(a, depth_a - depth_b)
end
if a == b
return a
end
(@size_bits - 1).downto(0) do |i|
if lca_ancestors[a][i] == -1 || lca_ancestors[b][i] == -1
next
end
if lca_ancestors[a][i] != lca_ancestors[b][i]
a = lca_ancestors[a][i]
b = lca_ancestors[b][i]
end
end
a = lca_ancestors[a][0]
b = lca_ancestors[b][0]
if a != b || a == -1 || b == -1
raise Exception.new("Assertion error")
end
a
end
def dist(a, b)
lca_node = lca(a, b)
lca_depths = @lca_depths.not_nil!
(lca_depths[a] - lca_depths[lca_node]) + (lca_depths[b] - lca_depths[lca_node])
end
# Returns `true` if node `c` is on the path from `a` to `b`.
def on_path?(a, b, c)
dist(a, c) + dist(c, b) == dist(a, b)
end
end
end
module AtCoder
# Implements [atcoder::mcf_graph](https://atcoder.github.io/ac-library/master/document_en/mincostflow.html).
#
# ```
# flow = AtCoder::MinCostFlow.new(5)
# flow.add_edge(0, 1, 30, 3)
# flow.add_edge(0, 2, 60, 9)
# flow.add_edge(1, 2, 40, 5)
# flow.add_edge(1, 3, 50, 7)
# flow.add_edge(2, 3, 20, 8)
# flow.add_edge(2, 4, 50, 6)
# flow.add_edge(3, 4, 60, 7)
# flow.flow(0, 4, 70) # => {70, 1080}
# ```
class MinCostFlow
private record EdgeInfo, capacity : Int64, cost : Int64 | Nil do
def self.zero
new(Int64::MAX, 0_i64)
end
def +(edge : EdgeInfo)
from_cost = @cost
to_cost = edge.cost
if from_cost.nil? || to_cost.nil?
return self.class.new(0_i64, nil)
end
self.class.new(min(@capacity, edge.capacity), from_cost + to_cost)
end
def >(edge : EdgeInfo)
from_dist = dist
to_dist = edge.dist
return true if from_dist.nil?
return false if to_dist.nil?
from_dist > to_dist
end
def dist
return nil if @capacity == 0
@cost
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
end
# Implements atcoder::mcf_graph g(n).
def initialize(@size : Int64)
@graph = AtCoder::DirectedGraph(Nil, EdgeInfo).new(@size)
@dists = Array(Int64 | Nil).new(@size, 0_i64)
@edges = [] of {Int64, Int64}
end
# Implements atcoder::mcf_graph.add_edge(from, to, capacity, cost).
def add_edge(from, to, capacity, cost)
@graph.add_edge(from, to, EdgeInfo.new(capacity.to_i64, cost.to_i64))
@graph.add_edge(to, from, EdgeInfo.new(0_i64, -cost.to_i64))
@edges << {from.to_i64, to.to_i64}
end
private def increment_edge_cost(from, to, cost_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, edge + EdgeInfo.new(edge.capacity, cost_diff))
end
private def increment_edge_capacity(from, to, capacity_diff)
edge = @graph.get_edge(from, to)
@graph.update_edge(from, to, EdgeInfo.new(edge.capacity + capacity_diff, edge.cost))
end
# Implements atcoder::mcf_graph.slope(start, target, flow_limit).
# ameba:disable Metrics/CyclomaticComplexity
def slope(start, target, flow_limit : Int | Nil = nil)
raise ArgumentError.new("start and target cannot be the same") if start == target
flow_points = [] of {Int64, Int64}
current_cost = 0_i64
flowed_capacity = 0_i64
min_cost = 0_i64
until flowed_capacity == flow_limit
nodes = @graph.dijkstra(start)
target_dist = nodes[target][:dist]
if target_dist.nil? || target_dist.capacity == 0
break
end
capacity = target_dist.capacity
# Update edge capacities on the path from start to target
last_node = target
until last_node == start
prev_node = nodes[last_node][:prev].not_nil!
increment_edge_capacity(prev_node, last_node, -capacity)
increment_edge_capacity(last_node, prev_node, capacity)
last_node = prev_node
end
# Update edge costs
@edges.each do |from, to|
from_dist = nodes[from][:dist]
to_dist = nodes[to][:dist]
next if from_dist.nil? || to_dist.nil?
next if from_dist.cost.nil? || to_dist.cost.nil?
dist = to_dist.cost.not_nil! - from_dist.cost.not_nil!
increment_edge_cost(from, to, -dist)
increment_edge_cost(to, from, dist)
end
# Update distants
nodes.each_with_index do |node, i|
dist = node[:dist]
if dist.nil? || @dists[i].nil? || dist.not_nil!.cost.nil?
@dists[i] = nil
else
@dists[i] = @dists[i].not_nil! + dist.not_nil!.cost.not_nil!
end
end
new_cost = @dists[target].not_nil!
if flow_limit.nil?
new_capacity = capacity
else
new_capacity = min(capacity, flow_limit - flowed_capacity)
end
if new_cost != current_cost
if current_cost == 0 && flowed_capacity != 0
flow_points << {0_i64, 0_i64}
end
flow_points << {flowed_capacity, min_cost}
end
min_cost += new_cost * new_capacity
flowed_capacity += new_capacity
current_cost = new_cost
end
flow_points << {flowed_capacity, min_cost}
flow_points
end
# Implements atcoder::mcf_graph.flow(start, target, flow_limit).
def flow(start, target, flow_limit : Int | Nil = nil)
flow_points = slope(start, target, flow_limit)
flow_points.last
end
@[AlwaysInline]
private def min(a, b)
a > b ? b : a
end
delegate :size, to: @graph
end
end
require "spec"
alias MinCostFlow = AtCoder::MinCostFlow
describe "MinCostFlow" do
describe ".flow" do
it "raises ArgumentError when start and target is the same" do
flow = MinCostFlow.new(2)
flow.add_edge(0, 1, 30, 3)
expect_raises(ArgumentError) { flow.flow(0, 0, 20) }
end
it "returns zero when there's no capacity anymore" do
flow = MinCostFlow.new(2)
flow.add_edge(1, 0, 30, 3)
flow.flow(0, 1, 0).should eq({0, 0})
end
it "calculates min-cost flow from start to target" do
flow = MinCostFlow.new(5)
flow.add_edge(0, 1, 30, 3)
flow.add_edge(0, 2, 60, 9)
flow.add_edge(1, 2, 40, 5)
flow.add_edge(1, 3, 50, 7)
flow.add_edge(2, 3, 20, 8)
flow.add_edge(2, 4, 50, 6)
flow.add_edge(3, 4, 60, 7)
flow.flow(0, 4, 70).should eq({70, 1080})
end
it "returns correct capacity and cost when there's multiple routes with the same costs" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 1)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.flow(0, 2, 20).should eq({20, 40})
end
it "returns slope to the maximum capacity if the third argument is omitted" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 2)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.flow(0, 2).should eq({20, 50})
end
it "returns slope to the maximum capacity if the capacity is full" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 2)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.flow(0, 2, 30).should eq({20, 50})
end
it "returns zero when the graph is not connected" do
flow = MinCostFlow.new(4)
flow.add_edge(0, 1, 30, 3)
flow.add_edge(2, 3, 30, 3)
flow.flow(0, 3, 10).should eq({0, 0})
end
end
describe ".slope" do
it "raises ArgumentError when start and target is the same" do
flow = MinCostFlow.new(2)
flow.add_edge(0, 1, 30, 3)
expect_raises(ArgumentError) { flow.slope(0, 0, 20) }
end
it "returns single point when there's no capacity anymore" do
flow = MinCostFlow.new(2)
flow.add_edge(1, 0, 30, 3)
flow.slope(0, 1, 0).should eq [
{0, 0},
]
end
it "calculates min-cost slope from start to target" do
flow = MinCostFlow.new(5)
flow.add_edge(0, 1, 30, 3)
flow.add_edge(0, 2, 60, 9)
flow.add_edge(1, 2, 40, 5)
flow.add_edge(1, 3, 50, 7)
flow.add_edge(2, 3, 20, 8)
flow.add_edge(2, 4, 50, 6)
flow.add_edge(3, 4, 60, 7)
flow.slope(0, 4, 70).should eq [
{0, 0},
{30, 420},
{50, 720},
{70, 1080},
]
end
it "returns line with slope with gradient zero if initial flow cost is zero" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 0)
flow.add_edge(1, 2, 10, 0)
flow.add_edge(0, 2, 10, 1)
flow.slope(0, 2, 15).should eq [
{0, 0},
{10, 0},
{15, 5},
]
end
it "returns single line when there's multiple routes with the same costs" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 1)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.slope(0, 2, 20).should eq [
{0, 0},
{20, 40},
]
end
it "returns slope to the maximum capacity if the third argument is omitted" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 2)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.slope(0, 2).should eq [
{0, 0},
{10, 20},
{20, 50},
]
end
it "returns slope to the maximum capacity if the capacity is full" do
flow = MinCostFlow.new(3)
flow.add_edge(0, 1, 10, 2)
flow.add_edge(1, 2, 10, 1)
flow.add_edge(0, 2, 10, 2)
flow.slope(0, 2, 30).should eq [
{0, 0},
{10, 20},
{20, 50},
]
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/mod_int.cr"
# require "./spec_helper.cr"
require "spec"
alias ModInt1000000007 = AtCoder::ModInt1000000007
alias ModInt7 = AtCoder::ModInt7
describe "ModInt1000000007" do
describe "#==" do
it "has a sound equality" do
ModInt1000000007.new(100_i64).should eq 100
ModInt1000000007.new(100_i64).should eq 100_i64
ModInt1000000007.new(100_i64).should_not eq 101
ModInt1000000007.new(100_i64).should_not eq 101_i64
ModInt1000000007.new(100_i64).should eq ModInt1000000007.new(100_i64)
ModInt1000000007.new(100_i64).should_not eq ModInt1000000007.new(101_i64)
(ModInt1000000007.new(100_i64) == 100).should eq true
(ModInt1000000007.new(100_i64) == 101).should eq false
(ModInt1000000007.new(100_i64) != 100).should eq false
(ModInt1000000007.new(100_i64) != 101).should eq true
end
end
describe ".new" do
it "constructs ModInt1000000007 from Int64" do
n = ModInt1000000007.new(100_i64)
n.should eq 100
end
it "rounds numbers larger than MOD" do
n = ModInt1000000007.new(1_000_000_007_i64)
n.should eq 0
n = ModInt1000000007.new(1_000_000_008_i64)
n.should eq 1
n = ModInt1000000007.new(3_000_000_025_i64)
n.should eq 4
n = ModInt1000000007.new(Int64::MAX)
n.should eq 291_172_003
end
it "rounds negative numbers" do
n = ModInt1000000007.new(-1_i64)
n.should eq 1_000_000_006
n = ModInt1000000007.new(-1_000_000_007_i64)
n.should eq 0
n = ModInt1000000007.new(Int64::MIN)
n.should eq 708_828_003
end
end
describe "#inv" do
it "generates inverse number under modulo" do
n = ModInt1000000007.new(1_i64)
n.inv.should eq n
n = ModInt1000000007.new(100_i64)
n.inv.should eq 570000004
n = ModInt1000000007.new(-1_i64 % 1_000_000_007_i64)
n.inv.should eq n
end
it "generates inverse number that multiplies with self into one" do
1_i64.step(by: 17, to: 1000) do |i|
(ModInt1000000007.new(i).inv * i).should eq 1
end
end
end
describe "#+" do
it "adds with number" do
n = ModInt1000000007.new(1_i64)
(n + 100).should eq 101
(n + 1_000_000_007_i64).should eq 1
end
end
describe "#-" do
it "subtracts with number" do
n = ModInt1000000007.new(100_i64)
(n - 10).should eq 90
(n - 1000).should eq -900 % 1_000_000_007_i64
end
end
describe "#*" do
it "multiplies with number" do
n = ModInt1000000007.new(100_i64)
(n * 334).should eq 33400
(n * -42).should eq -4200 % 1_000_000_007_i64
(n * 10_000_000_000_i64).should eq -7000 % 1_000_000_007_i64
end
end
describe "#/" do
it "divides with number" do
n = ModInt1000000007.new(700_i64)
(n // 5).should eq 140
(n // -1).should eq -700 % 1_000_000_007_i64
(n // 10_000_000_000_i64).should eq -10 % 1_000_000_007_i64
end
it "throws DivisionByZeroError when it tried to divide by zero" do
expect_raises(DivisionByZeroError) do
n = ModInt1000000007.new(700_i64)
n // 0
end
end
it "generates number that multiplies with self into the original one" do
1_i64.step(by: 9_837, to: 100_000) do |i|
1_000_000_i64.step(by: 9_837, to: 100_000) do |j|
(ModInt1000000007.new(i) // j * j).should eq i
end
end
end
end
describe "#**" do
it "generates powered number" do
n = ModInt1000000007.new(5_i64)
(n ** 0).should eq 1
(n ** 5).should eq 3125
(n ** -3).should eq ModInt1000000007.new(1_i64) // 125
(n ** 100_000).should eq 754_573_817
end
end
describe "#sqrt" do
it "generates square-rooted number" do
ModInt1000000007.new(100_i64).sqrt.should eq 10
ModInt1000000007.new(2_i64).sqrt.should eq 59_713_600
ModInt1000000007.new(5_i64).sqrt.should eq nil
end
end
describe "#to_s" do
it "serializes value as number" do
ModInt1000000007.new(100_i64).to_s.should eq "100"
ModInt1000000007.new(123_456_789_i64).to_s.should eq "123456789"
end
end
describe "#inspect" do
it "serializes value as number" do
ModInt1000000007.new(100_i64).inspect.should eq "100"
ModInt1000000007.new(123_456_789_i64).inspect.should eq "123456789"
end
end
describe ".factorial" do
it "calculates factorial fast" do
ModInt1000000007.factorial(3_i64).should eq 6
ModInt1000000007.factorial(10_i64).should eq 3628800
end
end
describe ".permutation" do
it "calculates permutation fast" do
ModInt1000000007.permutation(4_i64, 2_i64).should eq 12
ModInt1000000007.permutation(10_i64, 3_i64).should eq 720
end
end
describe ".combination" do
it "calculates combination fast" do
ModInt1000000007.combination(4_i64, 2_i64).should eq 6
ModInt1000000007.combination(10_i64, 3_i64).should eq 120
end
end
describe "Int" do
describe "#+" do
it "calculates addition with ModInt1000000007" do
(100 + ModInt1000000007.new(10_i64)).should eq 110
end
end
describe "#-" do
it "calculates substriction with ModInt1000000007" do
(100 - ModInt1000000007.new(10_i64)).should eq 90
end
end
describe "#*" do
it "calculates multiplication with ModInt1000000007" do
(100 * ModInt1000000007.new(10_i64)).should eq 1000
end
end
describe "#//" do
it "calculates multiplication with ModInt1000000007" do
(100 // ModInt1000000007.new(10_i64)).should eq 10
end
end
describe "#==" do
it "determines identity with ModInt1000000007" do
(100 == ModInt1000000007.new(100_i64)).should eq true
end
end
end
end
describe "static_modint" do
it "can define new ModInt record using custom modulo" do
n = ModInt7.new(3)
(n + 10).should eq 6
(n - 20).should eq 4
(n * 10).should eq 2
(n / 11).should eq 6
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/prime.cr"
require "spec"
alias Prime = AtCoder::Prime
describe "Prime" do
it "generates valid array of primes from the beginning" do
# https://oeis.org/A000040
expected = [
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59,
61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131,
137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197,
199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271,
]
Prime.first(58).should eq expected
end
describe ".prime_division" do
it "should factorize given number" do
Prime.prime_division(1).should eq [] of Int64
Prime.prime_division(175_i16).should eq [{5, 2}, {7, 1}]
Prime.prime_division(4096).should eq [{2, 12}]
Prime.prime_division(59049).should eq [{3, 10}]
Prime.prime_division(9_219_364_293_862_490_881_i64).should eq [{55103, 4}]
Prime.prime_division(4_611_686_018_427_387_904_i64).should eq [{2, 62}]
Prime.prime_division(1_000_000_000_000_000_000_i64).should eq [{2, 18}, {5, 18}]
Prime.prime_division(-1_000_000_i64).should eq [{-1, 1}, {2, 6}, {5, 6}]
Prime.prime_division(1_231_230_i64).should eq [{2, 1}, {3, 1}, {5, 1}, {7, 1}, {11, 1}, {13, 1}, {41, 1}]
Prime.prime_division(4_295_098_369_i64).should eq [{65537, 2}]
end
end
describe ".int_from_prime_division" do
it "should calculate weighted product of given base factors" do
Prime.int_from_prime_division([{2, 6}, {5, 6}]).should eq 1_000_000
Prime.int_from_prime_division([{2, 1}, {3, 1}, {5, 1}, {7, 1}, {11, 1}, {13, 1}, {41, 1}]).should eq 1_231_230
Prime.int_from_prime_division([{2_i64, 32_i64}]).should eq 4_294_967_296_i64
Prime.int_from_prime_division([{65521_u64, 4_u64}]).should eq 18_429_861_372_428_076_481_u64
end
end
describe ".prime?" do
it "is consistent with the result of prime generation" do
primes_under1000 = Prime.take_while { |prime| prime < 1000 }.to_set
0_i64.upto(1000_i64) do |n|
Prime.prime?(n).should eq primes_under1000.includes?(n)
end
end
it "can determine the primirity of large numbers" do
# https://en.wikipedia.org/wiki/Repunit#Factorization_of_decimal_repunits
Prime.prime?(1_i64).should be_false
Prime.prime?(11_i64).should be_true
Prime.prime?(111_i64).should be_false
Prime.prime?(1_111_i64).should be_false
Prime.prime?(11_111_i64).should be_false
Prime.prime?(111_111_i64).should be_false
Prime.prime?(1_111_111_i64).should be_false
Prime.prime?(11_111_111_i64).should be_false
Prime.prime?(111_111_111_i64).should be_false
Prime.prime?(1_111_111_111_i64).should be_false
Prime.prime?(11_111_111_111_i64).should be_false
Prime.prime?(111_111_111_111_i64).should be_false
Prime.prime?(1_111_111_111_111_i64).should be_false
Prime.prime?(11_111_111_111_111_i64).should be_false
Prime.prime?(111_111_111_111_111_i64).should be_false
Prime.prime?(1_111_111_111_111_111_i64).should be_false
Prime.prime?(11_111_111_111_111_111_i64).should be_false
Prime.prime?(111_111_111_111_111_111_i64).should be_false
Prime.prime?(1_111_111_111_111_111_111_i64).should be_true
# https://oeis.org/A123568
Prime.prime?(31_i64).should be_true
Prime.prime?(331_i64).should be_true
Prime.prime?(3_331_i64).should be_true
Prime.prime?(33_331_i64).should be_true
Prime.prime?(333_331_i64).should be_true
Prime.prime?(3_333_331_i64).should be_true
Prime.prime?(33_333_331_i64).should be_true
Prime.prime?(333_333_331_i64).should be_false
Prime.prime?(3_333_333_331_i64).should be_false
Prime.prime?(33_333_333_331_i64).should be_false
Prime.prime?(333_333_333_331_i64).should be_false
Prime.prime?(3_333_333_333_331_i64).should be_false
Prime.prime?(33_333_333_333_331_i64).should be_false
Prime.prime?(333_333_333_333_331_i64).should be_false
Prime.prime?(3_333_333_333_333_331_i64).should be_false
Prime.prime?(33_333_333_333_333_331_i64).should be_false
Prime.prime?(333_333_333_333_333_331_i64).should be_true
Prime.prime?(3_333_333_333_333_333_331_i64).should be_false
# https://en.wikipedia.org/wiki/Truncatable_prime#Decimal_truncatable_primes
Prime.prime?(7_i64).should be_true
Prime.prime?(37_i64).should be_true
Prime.prime?(137_i64).should be_true
Prime.prime?(9_137_i64).should be_true
Prime.prime?(29_137_i64).should be_true
Prime.prime?(629_137_i64).should be_true
Prime.prime?(7_629_137_i64).should be_true
Prime.prime?(67_629_137_i64).should be_true
Prime.prime?(567_629_137_i64).should be_true
Prime.prime?(6_567_629_137_i64).should be_true
Prime.prime?(16_567_629_137_i64).should be_true
Prime.prime?(216_567_629_137_i64).should be_true
Prime.prime?(6_216_567_629_137_i64).should be_true
Prime.prime?(46_216_567_629_137_i64).should be_true
Prime.prime?(646_216_567_629_137_i64).should be_true
Prime.prime?(2_646_216_567_629_137_i64).should be_true
Prime.prime?(12_646_216_567_629_137_i64).should be_true
Prime.prime?(312_646_216_567_629_137_i64).should be_true
end
end
describe ".each_divisor" do
it "should iterate through all positive divisor of given number" do
Prime.each_divisor(1).to_set.should eq Set{1}
Prime.each_divisor(100_i8).to_set.should eq Set{1, 2, 4, 5, 10, 20, 25, 50, 100}
Prime.each_divisor(4_295_098_369_i64).to_set.should eq Set{1, 65537, 4_295_098_369}
Prime.each_divisor(321_593_407_609_i64).to_set.should eq Set{1, 41, 7843741649, 321593407609}
end
it "raises if the argument is non-positive" do
expect_raises(ArgumentError) { Prime.each_divisor(0) }
expect_raises(ArgumentError) { Prime.each_divisor(-100_i64) }
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/priority_queue.cr"
require "spec"
alias PriorityQueue = AtCoder::PriorityQueue
describe "PriorityQueue" do
it "pops values in priority order" do
q = PriorityQueue(Int32).new
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 8
q.pop.should eq 7
q.pop.should eq 6
q.pop.should eq 5
q.pop.should eq 4
q.pop.should eq 3
q.pop.should eq 2
q.pop.should eq 1
q.pop.should eq nil
end
describe ".pop!" do
it "pops values in priority order. Raises Enumerable::NotFoundError if queue is empty." do
q = PriorityQueue(Int32).max
q << 5 << 6 << 1 << 3 << 2 << 8 << 7 << 4
q.pop!.should eq 8
q.pop!.should eq 7
q.pop!.should eq 6
q.pop!.should eq 5
q.pop!.should eq 4
q.pop!.should eq 3
q.pop!.should eq 2
q.pop!.should eq 1
expect_raises(Enumerable::EmptyError) { q.pop! }
end
end
it "initialized by elements in *enumerable*" do
q = PriorityQueue.max([5, 6, 1, 3, 2, 8, 7, 4])
q.pop.should eq 8
q.pop.should eq 7
q.pop.should eq 6
q.pop.should eq 5
q.pop.should eq 4
q.pop.should eq 3
q.pop.should eq 2
q.pop.should eq 1
q.pop.should eq nil
end
describe "with custom compare function" do
it "pops values in priority order" do
q = PriorityQueue(Int32).new { |a, b| a >= b }
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 1
q.pop.should eq 2
q.pop.should eq 3
q.pop.should eq 4
q.pop.should eq 5
q.pop.should eq 6
q.pop.should eq 7
q.pop.should eq 8
q.pop.should eq nil
end
end
describe ".max" do
it "pops values in ascending order of priority" do
q = PriorityQueue(Int32).max
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 8
q.pop.should eq 7
q.pop.should eq 6
q.pop.should eq 5
q.pop.should eq 4
q.pop.should eq 3
q.pop.should eq 2
q.pop.should eq 1
q.pop.should eq nil
end
end
describe ".min" do
it "pops values in descending order of priority" do
q = PriorityQueue(Int32).min
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 1
q.pop.should eq 2
q.pop.should eq 3
q.pop.should eq 4
q.pop.should eq 5
q.pop.should eq 6
q.pop.should eq 7
q.pop.should eq 8
q.pop.should eq nil
end
end
describe "#each" do
it "yields each item in the queue in comparator's order." do
array = [5, 6, 1, 3, 2, 8, 7, 4]
sorted = array.sort.reverse!
q = PriorityQueue.new(array)
i = 0
q.each do |x|
x.should eq sorted[i]
i += 1
end
end
end
describe "#first" do
it "returns, but does not remove, the head of the queue." do
q = PriorityQueue(Int32).new
q << 5
q << 6
q << 1
q << 3
q.first.should eq 6
q.first.should eq 6
q.pop.should eq 6
q.first.should eq 5
q.pop.should eq 5
q.first.should eq 3
q.pop.should eq 3
q.first.should eq 1
q.pop.should eq 1
q.first?.should eq nil
q.pop.should eq nil
end
end
describe "#empty?" do
it "should report true if queue is empty" do
q = PriorityQueue(Int32).new
q.empty?.should eq true
q << 1
q.empty?.should eq false
q.pop
q.empty?.should eq true
end
end
describe "#size" do
it "returns size of the queue" do
q = PriorityQueue(Int32).new
q << 1
q << 1
q << 1
q << 1
q.size.should eq 4
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/red_black_tree.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# bench_red_black_tree.rb by Rubinius Project, (c) Rubinius Contributors
# https://github.com/rubinius/rubinius-benchmark/blob/master/real_world/bench_red_black_tree.rb
# Licensed under BSD License https://spdx.org/licenses/BSD-3-Clause.html
# red_black_tree.cr by Crystal Language Project, (c) Manas Technology Solutions
# https://github.com/crystal-lang/crystal/blob/master/samples/red_black_tree.cr
# Licensed under MIT License https://mit-license.org
module AtCoder
class RedBlackTree
class Node
property :color
property :key
property! :left
property! :right
property! parent : self
RED = :red
BLACK = :black
def initialize(@key : Int32, @color = RED)
@left = @right = @parent = NilNode.instance
end
def black?
color == BLACK
end
def red?
color == RED
end
def nil_node?
false
end
end
class NilNode < Node
def self.instance
@@instance ||= RedBlackTree::NilNode.new
end
def initialize
@key = 0
@color = BLACK
@left = @right = @parent = self
end
def nil_node?
true
end
end
include Enumerable(Int32)
property root : Node
property :size
def initialize
@root = NilNode.instance
@size = 0
end
def insert(key)
insert_node(Node.new(key))
end
def insert_node(x)
insert_helper(x)
x.color = Node::RED
while x != root && x.parent.color == Node::RED
if x.parent == x.parent.parent.left
y = x.parent.parent.right
if !y.nil_node? && y.color == Node::RED
x.parent.color = Node::BLACK
y.color = Node::BLACK
x.parent.parent.color = Node::RED
x = x.parent.parent
else
if x == x.parent.right
x = x.parent
left_rotate(x)
end
x.parent.color = Node::BLACK
x.parent.parent.color = Node::RED
right_rotate(x.parent.parent)
end
else
y = x.parent.parent.left
if !y.nil_node? && y.color == Node::RED
x.parent.color = Node::BLACK
y.color = Node::BLACK
x.parent.parent.color = Node::RED
x = x.parent.parent
else
if x == x.parent.left
x = x.parent
right_rotate(x)
end
x.parent.color = Node::BLACK
x.parent.parent.color = Node::RED
left_rotate(x.parent.parent)
end
end
end
root.color = Node::BLACK
end
def delete_node(z)
y = (z.left.nil_node? || z.right.nil_node?) ? z : successor(z)
x = y.left.nil_node? ? y.right : y.left
x.parent = y.parent
if y.parent.nil_node?
self.root = x
else
if y == y.parent.left
y.parent.left = x
else
y.parent.right = x
end
end
z.key = y.key if y != z
if y.color == Node::BLACK
delete_fixup(x)
end
self.size -= 1
y
end
def minimum_node(x = root)
while !x.left.nil_node?
x = x.left
end
x
end
def maximum_node(x = root)
while !x.right.nil_node?
x = x.right
end
x
end
def successor(x)
if !x.right.nil_node?
return minimum_node(x.right)
end
y = x.parent
while !y.nil_node? && x == y.right
x = y
y = y.parent
end
y
end
def predecessor(x)
if !x.left.nil_node?
return maximum_node(x.left)
end
y = x.parent
while !y.nil_node? && x == y.left
x = y
y = y.parent
end
y
end
def inorder_walk(x = root, &)
# ameba:disable Lint/ShadowedArgument
x = self.minimum_node
while !x.nil_node?
yield x.key
x = successor(x)
end
end
def each(x = root, &)
inorder_walk(x) { |k| yield k }
end
def reverse_inorder_walk(x = root, &)
# ameba:disable Lint/ShadowedArgument
x = self.maximum_node
while !x.nil_node?
yield x.key
x = predecessor(x)
end
end
def reverse_each(x = root, &)
reverse_inorder_walk(x) { |k| yield k }
end
def search(key, x = root)
while !x.nil_node? && x.key != key
x = (key < x.key) ? x.left : x.right
end
x
end
def lower_than(key, x = root)
while !x.nil_node? && x.key != key
tmp = (key < x.key) ? x.left : x.right
if tmp.nil_node?
return x
end
x = tmp
end
x
end
def empty?
self.root.nil_node?
end
def black_height(x = root)
height = 0
while !x.nil_node?
x = x.left
height += 1 if x.nil_node? || x.black?
end
height
end
private def left_rotate(x)
raise "x.right is nil!" if x.right.nil_node?
y = x.right
x.right = y.left
y.left.parent = x if !y.left.nil_node?
y.parent = x.parent
if x.parent.nil_node?
self.root = y
else
if x == x.parent.left
x.parent.left = y
else
x.parent.right = y
end
end
y.left = x
x.parent = y
end
private def right_rotate(x)
raise "x.left is nil!" if x.left.nil_node?
y = x.left
x.left = y.right
y.right.parent = x if !y.right.nil_node?
y.parent = x.parent
if x.parent.nil_node?
self.root = y
else
if x == x.parent.left
x.parent.left = y
else
x.parent.right = y
end
end
y.right = x
x.parent = y
end
private def insert_helper(z)
y = NilNode.instance
x = root
while !x.nil_node?
y = x
x = (z.key < x.key) ? x.left : x.right
end
z.parent = y
if y.nil_node?
self.root = z
else
z.key < y.key ? y.left = z : y.right = z
end
self.size += 1
end
# ameba:disable Metrics/CyclomaticComplexity
private def delete_fixup(x)
while x != root && x.color == Node::BLACK
if x == x.parent.left
w = x.parent.right
if w.color == Node::RED
w.color = Node::BLACK
x.parent.color = Node::RED
left_rotate(x.parent)
w = x.parent.right
end
if w.left.color == Node::BLACK && w.right.color == Node::BLACK
w.color = Node::RED
x = x.parent
else
if w.right.color == Node::BLACK
w.left.color = Node::BLACK
w.color = Node::RED
right_rotate(w)
w = x.parent.right
end
w.color = x.parent.color
x.parent.color = Node::BLACK
w.right.color = Node::BLACK
left_rotate(x.parent)
x = root
end
else
w = x.parent.left
if w.color == Node::RED
w.color = Node::BLACK
x.parent.color = Node::RED
right_rotate(x.parent)
w = x.parent.left
end
if w.right.color == Node::BLACK && w.left.color == Node::BLACK
w.color = Node::RED
x = x.parent
else
if w.left.color == Node::BLACK
w.right.color = Node::BLACK
w.color = Node::RED
left_rotate(w)
w = x.parent.left
end
w.color = x.parent.color
x.parent.color = Node::BLACK
w.left.color = Node::BLACK
right_rotate(x.parent)
x = root
end
end
end
x.color = Node::BLACK
end
def min
minimum_node.key
end
def max
maximum_node.key
end
def delete(key)
node = search(key)
unless node.nil_node?
delete_node(node)
end
end
def has_key?(key)
node = search(key)
!node.nil_node?
end
def <<(x)
insert(x)
end
end
end
require "spec"
alias RedBlackTree = AtCoder::RedBlackTree
describe "RedBlackTree" do
it "pops values in priority order" do
tree = RedBlackTree.new
tree << 5
tree << 6
tree << 1
tree << 3
tree << 2
tree << 8
tree << 7
tree << 4
tree.max.should eq 8
tree.min.should eq 1
tree.delete 1
tree.delete 5
tree.delete 8
tree.max.should eq 7
tree.min.should eq 2
end
it "implements enumerable" do
tree = RedBlackTree.new
tree << 5
tree << 6
tree << 1
tree << 3
tree << 2
tree << 8
tree << 7
tree << 4
tree.first(5).should eq [1, 2, 3, 4, 5]
end
describe "#empty?" do
it "should report true if queue is empty" do
tree = RedBlackTree.new
tree.empty?.should eq true
tree << 1
tree.empty?.should eq false
tree.delete 1
tree.empty?.should eq true
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/seg_tree.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
# Implements [atcoder::segtree](https://atcoder.github.io/ac-library/master/document_en/segtree.html).
#
# The identity element will be implicitly defined as nil, so you don't
# have to manually define it. In the other words, you cannot include
# nil into an element of the monoid.
#
# ```
# tree = AtCoder::SegTree.new((0...100).to_a) { |a, b| [a, b].min }
# tree[10...50] # => 10
# ```
class SegTree(T)
getter values : Array(T)
@height : Int32
@n_leaves : Int32
def initialize(values : Array(T))
initialize(values) { |a, b| a > b ? a : b }
end
def initialize(@values : Array(T), &@operator : T, T -> T)
@height = log2_ceil(@values.size)
@n_leaves = 1 << @height
@segments = Array(T | Nil).new(2 * @n_leaves, nil)
# initialize segments
values.each_with_index { |x, i| @segments[@n_leaves + i] = x.as(T | Nil) }
(@n_leaves - 1).downto(1) { |i| refresh(i) }
end
@[AlwaysInline]
private def operate(a : T | Nil, b : T | Nil)
if a.nil?
b
elsif b.nil?
a
else
@operator.call(a, b)
end
end
# Implements atcoder::segtree.set(index, value)
def []=(index : Int, value : T)
@values[index] = value
index += @n_leaves
@segments[index] = value.as(T | Nil)
(1..@height).each { |j| refresh(ancestor(index, j)) }
end
# Implements atcoder::segtree.get(index)
def [](index : Int)
@values[index]
end
# Implements atcoder::segtree.prod(l, r)
def [](range : Range)
l = (range.begin || 0) + @n_leaves
r = (range.exclusive? ? (range.end || @n_leaves) : (range.end || @n_leaves - 1) + 1) + @n_leaves
sml, smr = nil.as(T | Nil), nil.as(T | Nil)
while l < r
if l.odd?
sml = operate(sml, @segments[l])
l += 1
end
if r.odd?
r -= 1
smr = operate(@segments[r], smr)
end
l >>= 1
r >>= 1
end
operate(sml, smr).not_nil!
end
# compatibility with ac-library
# Implements atcoder::segtree.set(index, value)
# alias of `.[]=`
def set(index : Int, value : T)
self.[]=(index, value)
end
# Implements atcoder::segtree.get(index)
# alias of `.[]`
def get(index : Int)
self.[](index)
end
# Implements atcoder::segtree.prod(left, right)
def prod(left : Int, right : Int)
self.[](left...right)
end
# Implements atcoder::segtree.all_prod(l, r)
def all_prod
@segments[1].not_nil!
end
# Implements atcoder::lazy_segtree.max_right(left, g).
def max_right(left, e : T | Nil = nil, & : T -> Bool)
unless 0 <= left && left <= @values.size
raise IndexError.new("{left: #{left}} must greater than or equal to 0 and less than or equal to {n: #{@values.size}}")
end
unless e.nil?
return nil unless yield e
end
return @values.size if left == @values.size
left += @n_leaves
sm = e
loop do
while left.even?
left >>= 1
end
res = operate(sm, @segments[left])
unless res.nil? || yield res
while left < @n_leaves
left = 2*left
res = operate(sm, @segments[left])
if res.nil? || yield res
sm = res
left += 1
end
end
return left - @n_leaves
end
sm = operate(sm, @segments[left])
left += 1
ffs = left & -left
break if ffs == left
end
@values.size
end
# Implements atcoder::lazy_segtree.min_left(right, g).
def min_left(right, e : T | Nil = nil, & : T -> Bool)
unless 0 <= right && right <= @values.size
raise IndexError.new("{right: #{right}} must greater than or equal to 0 and less than or equal to {n: #{@values.size}}")
end
unless e.nil?
return nil unless yield e
end
return 0 if right == 0
right += @n_leaves
sm = e
loop do
right -= 1
while right > 1 && right.odd?
right >>= 1
end
res = operate(@segments[right], sm)
unless res.nil? || yield res
while right < @n_leaves
right = 2*right + 1
res = operate(@segments[right], sm)
if res.nil? || yield res
sm = res
right -= 1
end
end
return right + 1 - @n_leaves
end
sm = operate(@segments[right], sm)
ffs = right & -right
break if ffs == right
end
0
end
@[AlwaysInline]
private def refresh(node : Int)
child1 = 2*node
child2 = 2*node + 1
@segments[node] = operate(@segments[child1], @segments[child2])
end
@[AlwaysInline]
private def ancestor(node, n_gens_ago)
node >> n_gens_ago
end
@[AlwaysInline]
private def log2_ceil(n : Int32) : Int32
sizeof(Int32)*8 - (n - 1).leading_zeros_count
end
@[AlwaysInline]
private def log2_ceil(n : Int32) : Int32
sizeof(Int32)*8 - (n - 1).leading_zeros_count
end
end
end
# require "../src/mod_int.cr"
# require "./spec_helper.cr"
require "spec"
alias SegTree = AtCoder::SegTree
describe "SegTree" do
describe "with default comparison" do
describe "#[]" do
it "simply returns value of tree" do
segtree = SegTree.new((1..100).to_a)
segtree[0].should eq 1
segtree[10].should eq 11
segtree[19].should eq 20
segtree[50].should eq 51
segtree[99].should eq 100
end
it "takes maximum value of selected range" do
segtree = SegTree.new((1..100).to_a)
segtree[0...100].should eq 100
segtree[10...50].should eq 50
segtree[50...70].should eq 70
segtree[12...33].should eq 33
segtree[0..0].should eq 1
segtree[16..16].should eq 17
segtree[49..49].should eq 50
segtree[88..88].should eq 89
segtree[10...].should eq 100
segtree[...60].should eq 60
segtree[...].should eq 100
segtree[10..].should eq 100
segtree[..60].should eq 61
segtree[..].should eq 100
end
end
describe "#[]=" do
it "updates value at selected index" do
segtree = SegTree.new((1..100).to_a)
segtree[40] = 9999
segtree[40].should eq 9999
segtree[0...100].should eq 9999
segtree[99] = 99999
segtree[99].should eq 99999
segtree[0...100].should eq 99999
end
end
end
describe "with custom comparison block" do
describe "min comparison" do
it "takes maximum value of selected range" do
segtree = SegTree.new((1..100).to_a) { |a, b| [a, b].min }
segtree[0...100].should eq 1
segtree[10...50].should eq 11
segtree[50...70].should eq 51
segtree[12...33].should eq 13
segtree[0..0].should eq 1
segtree[16..16].should eq 17
segtree[49..49].should eq 50
segtree[88..88].should eq 89
end
end
end
describe "#max_right" do
it "applies a binary search on the segment tree" do
segtree = SegTree(Int32).new((-10..10).map(&.**(2))) { |a, b| Math.min(a, b) }
(0..20).map { |i| segtree[0..i] }.to_a.should eq [100, 81, 64, 49, 36, 25, 16, 9, 4, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
segtree.max_right(0) { |x| x >= 10**9 }.should eq 0
segtree.max_right(0) { |x| x >= 20 }.should eq 6
segtree.max_right(0) { |x| x >= 16 }.should eq 7
segtree.max_right(0) { |x| x >= -1 }.should eq 21
segtree.max_right(0) { |x| x <= 10**9 }.should eq 21
segtree.max_right(0, e: Int32::MAX) { |x| x >= 10**9 }.should eq 0
segtree.max_right(0, e: Int32::MAX) { |x| x <= 10**9 }.should eq nil
end
end
describe "#min_left" do
it "applies a binary search on the segment tree" do
segtree = SegTree(Int32).new((-10..10).map(&.**(2))) { |a, b| Math.min(a, b) }
(0..20).map { |i| segtree[i..20] }.to_a.should eq [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
segtree.min_left(21) { |x| x >= 10**9 }.should eq 21
segtree.min_left(21) { |x| x >= 20 }.should eq 15
segtree.min_left(21) { |x| x >= 16 }.should eq 14
segtree.min_left(21) { |x| x >= -1 }.should eq 0
segtree.min_left(21) { |x| x <= 10**9 }.should eq 0
segtree.min_left(21, e: Int32::MAX) { |x| x <= 10**9 }.should eq nil
segtree.min_left(21, e: Int32::MAX) { |x| x >= 10**9 }.should eq 21
end
end
it "can be used with ModInt" do
values = 100.times.map { |i| Mint.new(1_i64) << i }.to_a
segtree = SegTree.new(values) { |a, b| a + b }
segtree[0...32].should eq 0xFFFFFFFF % Mint::MOD
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/skew_heap.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
class SkewHeap(T)
@root : Node(T) | Nil
@size : Int64
getter :size
class Node(T)
@left : Node(T) | Nil
@right : Node(T) | Nil
@value : T
property :value, :left, :right
def initialize(@left, @right, @value)
end
end
def initialize
initialize { |a, b| a <= b }
end
def initialize(&block : T, T -> Bool)
@root = nil
@size = 0_i64
@compare_proc = block
end
def meld(a : Node(T) | Nil, b : Node(T) | Nil)
return b if a.nil?
return a if b.nil?
if @compare_proc.call(a.value, b.value)
a, b = b, a
end
a.right = meld(a.right, b)
a.left, a.right = a.right, a.left
a
end
def push(value : T)
@root = meld(@root, Node(T).new(nil, nil, value))
@size += 1
end
def pop
return nil if @root.nil?
root = @root.not_nil!
ret = root.value
@root = meld(root.left, root.right)
@size -= 1
ret
end
# Alias of `push`
def <<(v : T)
push(v)
end
def empty?
@size == 0
end
end
end
require "spec"
alias SkewHeap = AtCoder::SkewHeap
describe "SkewHeap" do
it "pops values in priority order" do
q = SkewHeap(Int32).new
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 8
q.pop.should eq 7
q.pop.should eq 6
q.pop.should eq 5
q.pop.should eq 4
q.pop.should eq 3
q.pop.should eq 2
q.pop.should eq 1
q.pop.should eq nil
end
describe "with custom compare function" do
it "pops values in priority order" do
q = SkewHeap(Int32).new { |a, b| a >= b }
q << 5
q << 6
q << 1
q << 3
q << 2
q << 8
q << 7
q << 4
q.pop.should eq 1
q.pop.should eq 2
q.pop.should eq 3
q.pop.should eq 4
q.pop.should eq 5
q.pop.should eq 6
q.pop.should eq 7
q.pop.should eq 8
q.pop.should eq nil
end
end
describe "#empty?" do
it "should report true if queue is empty" do
q = SkewHeap(Int32).new
q.empty?.should eq true
q << 1
q.empty?.should eq false
q.pop
q.empty?.should eq true
end
end
describe "#size" do
it "returns size of the queue" do
q = SkewHeap(Int32).new
q << 1
q << 1
q << 1
q << 1
q.size.should eq 4
end
end
end
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# require "../src/string.cr"
# ac-library.cr by hakatashi https://github.com/google/ac-library.cr
#
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
module AtCoder
module String
private SA_THRESHOLD_NAIVE = 10
private SA_THRESHOLD_DOUBLING = 40
private def self.sa_naive(s)
n = s.size
sa = (0...n).to_a.sort { |l, r|
next 1 if l == r
res = nil
while l < n && r < n
if s[l] != s[r]
res = s[l] <=> s[r]
break
end
l += 1
r += 1
end
res ? res : (l == n ? -1 : 1)
}
sa
end
private def self.sa_doubling(s)
n = s.size
sa = (0...n).to_a
rank = s.clone
tmp = [0] * n
k = 1
while k < n
cmp = ->(x : Int32, y : Int32) {
return rank[x] <=> rank[y] if rank[x] != rank[y]
rx = x + k < n ? rank[x + k] : -1
ry = y + k < n ? rank[y + k] : -1
rx <=> ry
}
sa.sort! { |a, b| cmp.call(a, b) }
tmp[sa[0]] = 0
(1...n).each do |i|
tmp[sa[i]] = tmp[sa[i - 1]] + (cmp.call(sa[i - 1], sa[i]) == -1 ? 1 : 0)
end
tmp, rank = rank, tmp
k *= 2
end
sa
end
# ameba:disable Metrics/CyclomaticComplexity
private def self.sa_is(s, upper)
n = s.size
case n
when .==(0)
return [] of Int32
when .==(1)
return [0]
when .==(2)
return s[0] < s[1] ? [0, 1] : [1, 0]
when .<(SA_THRESHOLD_NAIVE)
return sa_naive(s)
when .<(SA_THRESHOLD_DOUBLING)
return sa_doubling(s)
end
sa = [0] * n
ls = [false] * n
(n - 2).downto(0) do |i|
ls[i] = s[i] == s[i + 1] ? ls[i + 1] : s[i] < s[i + 1]
end
sum_l = [0] * (upper + 1)
sum_s = [0] * (upper + 1)
n.times do |i|
if ls[i]
sum_l[s[i] + 1] += 1
else
sum_s[s[i]] += 1
end
end
(0..upper).each do |i|
sum_s[i] += sum_l[i]
sum_l[i + 1] += sum_s[i] if i < upper
end
induce = ->(lms : Array(Int32)) do
sa = [-1] * n
buffer = sum_s.clone
lms.each do |d|
next if d == n
sa[buffer[s[d]]] = d
buffer[s[d]] += 1
end
buffer = sum_l.clone
sa[buffer[s[n - 1]]] = n - 1
buffer[s[n - 1]] += 1
n.times do |i|
v = sa[i]
if v >= 1 && !ls[v - 1]
sa[buffer[s[v - 1]]] = v - 1
buffer[s[v - 1]] += 1
end
end
buffer = sum_l.clone
(n - 1).downto(0) do |i|
v = sa[i]
if v >= 1 && ls[v - 1]
buffer[s[v - 1] + 1] -= 1
sa[buffer[s[v - 1] + 1]] = v - 1
end
end
end
lms_map = [-1] * (n + 1)
m = 0
(1...n).each do |i|
if !ls[i - 1] && ls[i]
lms_map[i] = m
m += 1
end
end
lms = Array(Int32).new(m)
(1...n).each do |i|
if !ls[i - 1] && ls[i]
lms << i
end
end
induce.call(lms)
return sa if m == 0
sorted_lms = Array(Int32).new(m)
sa.each do |v|
sorted_lms << v if lms_map[v] != -1
end
rec_s = [0] * m
rec_upper = 0
rec_s[lms_map[sorted_lms[0]]] = 0
(1...m).each do |i|
l = sorted_lms[i - 1]
r = sorted_lms[i]
end_l = (lms_map[l] + 1 < m) ? lms[lms_map[l] + 1] : n
end_r = (lms_map[r] + 1 < m) ? lms[lms_map[r] + 1] : n
same = true
if end_l - l != end_r - r
same = false
else
while l < end_l
break if s[l] != s[r]
l += 1
r += 1
end
same = false if l == n || s[l] != s[r]
end
rec_upper += 1 unless same
rec_s[lms_map[sorted_lms[i]]] = rec_upper
end
rec_sa = sa_is(rec_s, rec_upper)
m.times do |i|
sorted_lms[i] = lms[rec_sa[i]]
end
induce.call(sorted_lms)
sa
end
# returns suffix array in O(n + upper)
def self.suffix_array(sequence : Indexable(Int32), upper)
sa_is(sequence, upper)
end
# returns suffix array in O(n log(n))
def self.suffix_array(sequence : Indexable)
n = sequence.size
indices = (0...n).to_a.sort { |l, r| sequence[l] <=> sequence[r] }
s2 = [0] * n
now = 0
n.times do |i|
now += 1 if i > 0 && sequence[indices[i - 1]] != sequence[indices[i]]
s2[indices[i]] = now
end
upper = now
sa_is(s2, upper)
end
# returns suffix array in O(n)
def self.suffix_array(sequence : ::String)
sa_is(sequence.bytes.map(&.to_i32), 255)
end
# returns lcp array in O(n)
def self.lcp_array(sequence, sa)
n = sequence.size
rank = [0] * n
sa.each_with_index { |e, i| rank[e] = i }
lcp = [0] * (n - 1)
h = 0
n.times do |i|
h -= 1 if h > 0
next if rank[i] == 0
j = sa[rank[i] - 1]
while j + h < n && i + h < n
break if sequence[j + h] != sequence[i + h]
h += 1
end
lcp[rank[i] - 1] = h
end
lcp
end
# returns z array
def self.z_algorithm(sequence)
n = sequence.size
return [] of Int32 if n == 0
z = [0] * n
i, j = 1, 0
while i < n
z[i] = (j + z[j] <= i) ? 0 : (j + z[j] - i < z[i - j] ? j + z[j] - i : z[i - j])
while i + z[i] < n && sequence[z[i]] == sequence[i + z[i]]
z[i] += 1
end
j = i if j + z[j] < i + z[i]
i += 1
end
z[0] = n
z
end
end
end
require "spec"
describe "String" do
describe ".suffix_array" do
it "returns suffix array in O(n)" do
AtCoder::String.suffix_array("abcbcba").should eq [6, 0, 5, 3, 1, 4, 2]
AtCoder::String.suffix_array("mississippi").should eq [10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2]
AtCoder::String.suffix_array("ababacaca").should eq [8, 0, 2, 6, 4, 1, 3, 7, 5]
AtCoder::String.suffix_array("aaaaa").should eq [4, 3, 2, 1, 0]
end
it "returns suffix array in O(n log(n))" do
AtCoder::String.suffix_array([1_i64, 2_i64, 3_i64, 2_i64, 3_i64, 2_i64, 1_i64]).should eq [6, 0, 5, 3, 1, 4, 2]
AtCoder::String.suffix_array([2**12, 2**8, 2**18, 2**18, 2**8, 2**18, 2**18, 2**8, 2**15, 2**15, 2**8]).should eq [10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2]
AtCoder::String.suffix_array([1, 2, 1, 2, 1, 3, 1, 3, 1]).should eq [8, 0, 2, 6, 4, 1, 3, 7, 5]
AtCoder::String.suffix_array([10_i64**18, 10_i64**18, 10_i64**18, 10_i64**18, 10_i64**18]).should eq [4, 3, 2, 1, 0]
end
it "returns suffix array in O(n + uppper)" do
AtCoder::String.suffix_array([1, 2, 3, 2, 3, 2, 1], 3).should eq [6, 0, 5, 3, 1, 4, 2]
AtCoder::String.suffix_array([2**12, 2**8, 2**18, 2**18, 2**8, 2**18, 2**18, 2**8, 2**15, 2**15, 2**8], 2**18).should eq [10, 7, 4, 1, 0, 9, 8, 6, 3, 5, 2]
AtCoder::String.suffix_array([1, 2, 1, 2, 1, 3, 1, 3, 1], 3).should eq [8, 0, 2, 6, 4, 1, 3, 7, 5]
AtCoder::String.suffix_array([0, 0, 0, 0, 0], 10**8).should eq [4, 3, 2, 1, 0]
end
end
describe ".lcp_array" do
it "returns lcp_array" do
AtCoder::String.lcp_array("abcbcba", AtCoder::String.suffix_array("abcbcba")).should eq [1, 0, 1, 3, 0, 2]
AtCoder::String.lcp_array("mississippi", AtCoder::String.suffix_array("mississippi")).should eq [1, 1, 4, 0, 0, 1, 0, 2, 1, 3]
AtCoder::String.lcp_array("ababacaca", AtCoder::String.suffix_array("ababacaca")).should eq [1, 3, 1, 3, 0, 2, 0, 2]
AtCoder::String.lcp_array("aaaaa", AtCoder::String.suffix_array("aaaaa")).should eq [1, 2, 3, 4]
AtCoder::String.lcp_array([1, 2, 3, 2, 3, 2, 1], AtCoder::String.suffix_array([1, 2, 3, 2, 3, 2, 1])).should eq [1, 0, 1, 3, 0, 2]
AtCoder::String.lcp_array([12, 8, 18, 18, 8, 18, 18, 8, 15, 15, 8], AtCoder::String.suffix_array([12, 8, 18, 18, 8, 18, 18, 8, 15, 15, 8])).should eq [1, 1, 4, 0, 0, 1, 0, 2, 1, 3]
AtCoder::String.lcp_array([1, 2, 1, 2, 1, 3, 1, 3, 1], AtCoder::String.suffix_array([1, 2, 1, 2, 1, 3, 1, 3, 1])).should eq [1, 3, 1, 3, 0, 2, 0, 2]
AtCoder::String.lcp_array([1, 1, 1, 1, 1], AtCoder::String.suffix_array([1, 1, 1, 1, 1])).should eq [1, 2, 3, 4]
end
end
describe ".z_algorithm" do
it "returns empty array if input sequence is empty" do
AtCoder::String.z_algorithm([] of Int32).should eq [] of Int32
AtCoder::String.z_algorithm("").should eq [] of Int32
end
it "returns the sequence of length n, such that the i-th element is the length of the LCP (Longest Common Prefix) of s[0...n] and s[i...n]." do
AtCoder::String.z_algorithm("abcbcba").should eq [7, 0, 0, 0, 0, 0, 1]
AtCoder::String.z_algorithm("mississippi").should eq [11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
AtCoder::String.z_algorithm("ababacaca").should eq [9, 0, 3, 0, 1, 0, 1, 0, 1]
AtCoder::String.z_algorithm("aaaaa").should eq [5, 4, 3, 2, 1]
AtCoder::String.z_algorithm([1, 2, 3, 2, 3, 2, 1]).should eq [7, 0, 0, 0, 0, 0, 1]
AtCoder::String.z_algorithm([1, 2, 3, 3, 2, 3, 3, 2, 4, 4, 2]).should eq [11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
AtCoder::String.z_algorithm([1, 2, 1, 2, 1, 3, 1, 3, 1]).should eq [9, 0, 3, 0, 1, 0, 1, 0, 1]
AtCoder::String.z_algorithm([1, 1, 1, 1, 1]).should eq [5, 4, 3, 2, 1]
end
end
end