Class: Convolution

Inherits:
Object
  • Object
show all
Defined in:
lib/convolution.rb

Overview

usage :

conv = Convolution.new(mod, [primitive_root]) conv.convolution(a, b) #=> convolution a and b modulo mod.

Instance Method Summary collapse

Constructor Details

#initialize(mod = 998_244_353, primitive_root = nil) ⇒ Convolution

Returns a new instance of Convolution.



7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
 
# File 'lib/convolution.rb', line 7

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

Instance Method Details

#convolution(a, b) ⇒ Object

Raises:

  • (ArgumentError) 


34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
 
# File 'lib/convolution.rb', line 34

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