class AcLibraryRb::Convolution
usage :
conv = Convolution.new
(mod, [primitive_root]) conv.convolution(a, b) #=> convolution a and b modulo mod.
Public Class Methods
new(mod = 998_244_353, primitive_root = nil)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 8 def initialize(mod = 998_244_353, primitive_root = nil) @mod = mod cnt2 = bsf(@mod - 1) e = (primitive_root || calc_primitive_root(mod)).pow((@mod - 1) >> cnt2, @mod) ie = e.pow(@mod - 2, @mod) es = [0] * (cnt2 - 1) ies = [0] * (cnt2 - 1) cnt2.downto(2){ |i| es[i - 2] = e ies[i - 2] = ie e = e * e % @mod ie = ie * ie % @mod } now = inow = 1 @sum_e = [0] * cnt2 @sum_ie = [0] * cnt2 (cnt2 - 1).times{ |i| @sum_e[i] = es[i] * now % @mod now = now * ies[i] % @mod @sum_ie[i] = ies[i] * inow % @mod inow = inow * es[i] % @mod } end
Public Instance Methods
convolution(a, b)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 35 def convolution(a, b) n = a.size m = b.size return [] if n == 0 || m == 0 h = (n + m - 2).bit_length raise ArgumentError if h > @sum_e.size z = 1 << h a = a + [0] * (z - n) b = b + [0] * (z - m) batterfly(a, h) batterfly(b, h) c = a.zip(b).map{ |a, b| a * b % @mod } batterfly_inv(c, h) iz = z.pow(@mod - 2, @mod) return c[0, n + m - 1].map{ |c| c * iz % @mod } end
Private Instance Methods
batterfly(a, h)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 59 def batterfly(a, h) 1.upto(h){ |ph| w = 1 << (ph - 1) p = 1 << (h - ph) now = 1 w.times{ |s| offset = s << (h - ph + 1) offset.upto(offset + p - 1){ |i| l = a[i] r = a[i + p] * now % @mod a[i] = l + r a[i + p] = l - r } now = now * @sum_e[bsf(~s)] % @mod } } end
batterfly_inv(a, h)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 77 def batterfly_inv(a, h) h.downto(1){ |ph| w = 1 << (ph - 1) p = 1 << (h - ph) inow = 1 w.times{ |s| offset = s << (h - ph + 1) offset.upto(offset + p - 1){ |i| l = a[i] r = a[i + p] a[i] = l + r a[i + p] = (l - r) * inow % @mod } inow = inow * @sum_ie[bsf(~s)] % @mod } } end
bsf(x)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 95 def bsf(x) (x & -x).bit_length - 1 end
calc_primitive_root(mod)
click to toggle source
# File lib_lock/ac-library-rb/convolution.rb, line 99 def calc_primitive_root(mod) return 1 if mod == 2 return 3 if mod == 998_244_353 divs = [2] x = (mod - 1) / 2 x /= 2 while x.even? i = 3 while i * i <= x if x % i == 0 divs << i x /= i while x % i == 0 end i += 2 end divs << x if x > 1 g = 2 loop{ return g if divs.none?{ |d| g.pow((mod - 1) / d, mod) == 1 } g += 1 } end