Skip to content

Xor convolution

Rust

 1
 2
 3
 4
 5
 6
 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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
mod math {
    use num_traits::Zero;
    type Int = u32;

    pub fn xor_convolution(a: &Vec<Int>, b: &Vec<Int>) -> Vec<Int> {
        let n = 1 << next_power_of_2(a.len().max(b.len()));
        let mut fa = a.clone();
        let mut fb = b.clone();
        fa.resize(n, Int::zero());
        fb.resize(n, Int::zero());

        fast_walsh_hadamard_transform(&mut fa);
        fast_walsh_hadamard_transform(&mut fb);
        for i in 0..n {
            fa[i] *= fb[i];
        }
        fast_walsh_hadamard_transform(&mut fa);
        for x in &mut fa {
            *x /= n as Int; // Int = modint の場合は逆元を掛ける
        }
        fa
    }

    pub fn fast_walsh_hadamard_transform(x: &mut Vec<Int>) {
        let k = x.len().trailing_zeros() as usize;
        for p in 0..k {
            let step = 1 << p;
            for i in (0..x.len()).step_by(step * 2) {
                for j in i..i + step {
                    let u = x[j];
                    let v = x[j + step];
                    x[j] = u + v;
                    x[j + step] = u - v;
                }
            }
        }
    }

    ///  x <= 2^k を満たす最小の非負整数 k を求める。
    fn next_power_of_2(mut x: usize) -> usize {
        if x == 0 {
            return 0;
        }
        let mut k = 0;
        while (1 << k) < x {
            k += 1;
        }
        k
    }
}