階乗 mod pの高速な実装 (Rust)

階乗 mod 素数 - memo のO(p^{1/2} log p)解法の、Rustでの実装を与えた。説明は元記事の方で十分にされているので、この記事での説明は期待しないでほしい。また、実装にはところどころ雑な部分があるので、今後refineするかもしれない。

問題

n! mod pを計算せよ。
No.502 階乗を計算するだけ - yukicoder

コードの仕様

MOD: i64: mod。今回は10^9 + 7

登場する典型

元記事にない補足事項

元記事における整数vは、この実装では2^15に決め打ちしている。これは、v * v >= p = 10^9 + 7を満たす最小の2ベキである。

提出: https://yukicoder.me/submissions/347989:#347989 (Rust)

// mod_int::ModInt 省略

macro_rules! define_mod {
    ($struct_name: ident, $modulo: expr) => {
        #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
        struct $struct_name {}
        impl mod_int::Mod for $struct_name { fn m() -> i64 { $modulo } }
    }
}
const MOD: i64 = 1_000_000_007;
define_mod!(P, MOD);
type ModInt = mod_int::ModInt<P>;

/// FFT (in-place, verified as NTT only)
/// R: Ring + Copy
/// Verified by: https://codeforces.com/contest/1096/submission/47672373
mod fft {
    use std::ops::*;
    /// n should be a power of 2. zeta is a primitive n-th root of unity.
    /// one is unity
    /// Note that the result should be multiplied by 1/sqrt(n).
    pub fn transform<R>(f: &mut [R], zeta: R, one: R)
        where R: Copy +
        Add<Output = R> +
        Sub<Output = R> +
        Mul<Output = R> {
        let n = f.len();
        assert!(n.is_power_of_two());
        {
            let mut i = 0;
            for j in 1 .. n - 1 {
                let mut k = n >> 1;
                loop {
                    i ^= k;
                    if k <= i { break; }
                    k >>= 1;
                }
                if j < i { f.swap(i, j); }
            }
        }
        let mut zetapow = Vec::new();
        {
            let mut m = 1;
            let mut cur = zeta;
            while m < n {
                zetapow.push(cur);
                cur = cur * cur;
                m *= 2;
            }
        }
        let mut m = 1;
        while m < n {
            let base = zetapow.pop().unwrap();
            let mut r = 0;
            while r < n {
                let mut w = one;
                for s in r .. r + m {
                    let u = f[s];
                    let d = f[s + m] * w;
                    f[s] = u + d;
                    f[s + m] = u - d;
                    w = w * base;
                }
                r += 2 * m;
            }
            m *= 2;
        }
    }
}

mod arbitrary_mod {
    use mod_int;
    use fft;
    const MOD1: i64 = 1012924417;
    const MOD2: i64 = 1224736769;
    const MOD3: i64 = 1007681537;
    const G1: i64 = 5;
    const G2: i64 = 3;
    const G3: i64 = 3;
    define_mod!(P1, MOD1);
    define_mod!(P2, MOD2);
    define_mod!(P3, MOD3);

    fn zmod(mut a: i64, b: i64) -> i64 {
        a %= b;
        if a < 0 {
            a += b;
        }
        a
    }
    fn ext_gcd(mut a: i64, mut b: i64) -> (i64, i64, i64) {
        let mut x = 0;
        let mut y = 1;
        let mut u = 1;
        let mut v = 0;
        while a != 0 {
            let q = b / a;
            x -= q * u;
            std::mem::swap(&mut x, &mut u);
            y -= q * v;
            std::mem::swap(&mut y, &mut v);
            b -= q * a;
            std::mem::swap(&mut b, &mut a);
        }
        (b, x, y)
    }
    fn invmod(a: i64, b: i64) -> i64 {
        let x = ext_gcd(a, b).1;
        zmod(x, b)
    }

    // This function is ported from http://math314.hateblo.jp/entry/2015/05/07/014908
    fn garner(mut mr: Vec<(i64, i64)>, mo: i64) -> i64 {
        mr.push((mo, 0));

        let mut coffs = vec![1; mr.len()];
        let mut constants = vec![0; mr.len()];
        for i in 0..mr.len() - 1 {
            let v = zmod(mr[i].1 - constants[i], mr[i].0) * invmod(coffs[i], mr[i].0) % mr[i].0;
            assert!(v >= 0);
            for j in i + 1..mr.len() {
                constants[j] += coffs[j] * v % mr[j].0;
                constants[j] %= mr[j].0;
                coffs[j] = coffs[j] * mr[i].0 % mr[j].0;
            }
        }
        constants[mr.len() - 1]
    }

    // f *= g, g is destroyed
    fn convolution_friendly<P: mod_int::Mod>(a: &[i64], b: &[i64], gen: i64) -> Vec<i64> {
        use mod_int::ModInt;
        let d = a.len();
        let mut f = vec![ModInt::<P>::new(0); d];
        let mut g = vec![ModInt::<P>::new(0); d];
        for i in 0..d {
            f[i] = a[i].into();
            g[i] = b[i].into();
        }
        let zeta = ModInt::new(gen).pow((P::m() - 1) / d as i64);
        fft::transform(&mut f, zeta, ModInt::new(1));
        fft::transform(&mut g, zeta, ModInt::new(1));
        for i in 0..d {
            f[i] *= g[i];
        }
        fft::transform(&mut f, zeta.inv(), ModInt::new(1));
        let inv = ModInt::new(d as i64).inv();
        let mut ans = vec![0; d];
        for i in 0..d {
            ans[i] = (f[i] * inv).x;
        }
        ans
    }


    pub fn arbmod_convolution(a: &mut [i64], b: &mut [i64], mo: i64)
                          -> Vec<i64> {
        use ::mod_int::Mod;
        let d = a.len();
        assert!(d.is_power_of_two());
        assert_eq!(d, b.len());
        for x in a.iter_mut() {
            *x = zmod(*x, mo);
        }
        for x in b.iter_mut() {
            *x = zmod(*x, mo);
        }
        let x = convolution_friendly::<P1>(&a, &b, G1);
        let y = convolution_friendly::<P2>(&a, &b, G2);
        let z = convolution_friendly::<P3>(&a, &b, G3);

        let mut ret = vec![0; d];
        let mut mr = [(0, 0); 3];
        for i in 0..d {
            mr[0] = (P1::m(), x[i]);
            mr[1] = (P2::m(), y[i]);
            mr[2] = (P3::m(), z[i]);
            ret[i] = garner(mr.to_vec(), mo);
        }
        ret
    }
}

// f *= g, g is not destroyed
fn convolution(f: &mut [i64], g: &mut [i64]) {
    let ans = arbitrary_mod::arbmod_convolution(f, g, MOD);
    for i in 0..f.len() {
        f[i] = ans[i];
    }
}

fn grow(d: i64, v: i64, mut h: Vec<i64>,
        invfac: &[ModInt]) -> Vec<i64> {
    assert_eq!(h.len() as i64, d + 1);
    let dd = d as usize;
    let dm = ModInt::new(d);
    let vm = ModInt::new(v);

    let mut aux = vec![1; dd];

    let mut f = vec![0; 4 * dd];
    let mut g = vec![0; 4 * dd];
    for i in 0..dd + 1 {
        f[i] = (invfac[i] * invfac[dd - i] * h[i]).x;
        if (dd + i) % 2 != 0 {
            f[i] = if f[i] == 0 { 0 } else { MOD - f[i] };
        }
    }
    let oldf = f.clone();
    for (idx, &a) in [dm + 1, dm * vm.inv(), dm * vm.inv() + dm + 1].iter().enumerate() {
        for i in 0..4 * dd { f[i] = oldf[i]; }
        for i in 0..4 * dd { g[i] = 0; }
        for i in 1..2 * dd + 2 {
            g[i] = (a - d + i as i64 - 1).inv().x;
        }
        convolution(&mut f, &mut g);
        let mut prod = 1;
        for i in 0..dd + 1 {
            prod = prod * (a - i as i64).x % MOD;
            assert_ne!(prod, 0);
        }
        for i in 0..dd + 1 {
            f[dd + i + 1] = f[dd + i + 1] * prod % MOD;
            prod = prod * (a + i as i64 + 1).x % MOD;
            prod = prod * (a - d + i as i64).inv().x % MOD;
        }
        match idx {
            1 => {
                for i in 0..dd + 1 {
                    h[i] = h[i] * f[dd + 1 + i] % MOD;
                }
            }
            0 => {
                for i in 0..dd {
                    aux[i] = f[dd + 1 + i];
                }
            }
            2 => {
                for i in 0..dd {
                    aux[i] = aux[i] * f[dd + 1 + i] % MOD;
                }
            }
            _ => unreachable!(),
        }
    }
    h.extend_from_slice(&aux);
    h
}

fn gen_seq(d: i64, v: i64) -> Vec<i64> {
    assert!(d > 0 && (d as u64).is_power_of_two());
    let dd = d as usize;

    // precompute factorial and its inv
    let mut fac = vec![ModInt::new(0); 2 * dd + 1];
    let mut invfac = vec![ModInt::new(0); 2 * dd + 1];
    fac[0] = ModInt::new(1);
    for i in 1..2 * dd + 1 {
        fac[i] = fac[i - 1] * (i as i64);
    }
    invfac[2 * dd] = fac[2 * dd].inv();
    for i in (0..2 * dd).rev() {
        invfac[i] = invfac[i + 1] * (i as i64 + 1);
    }
    let mut size = 1;
    // Initialized with [g_1(0), g_1(v)].
    let mut seq = vec![1.into(), (v + 1).into()];
    while size < d {
        seq = grow(size, v, seq, &invfac);
        size *= 2;
    }
    assert_eq!(size, d);
    seq
}

fn fact(n: i64) -> ModInt {
    let d = 1 << 15;
    let aux = gen_seq(d, d);
    // eprintln!("{:?}", aux);
    let mut ans = ModInt::new(1);
    let lim = min(d, (n + 1) / d);
    for i in 0..lim {
        ans *= aux[i as usize];
    }
    for i in lim * d..n {
        ans *= i + 1;
    }
    ans
}

// Uses techniques described in https://min-25.hatenablog.com/entry/2017/04/10/215046.
// Bostan, A., Gaudry, P., & Schost, É. (2007). Linear Recurrences with Polynomial Coefficients and Application to Integer Factorization and Cartier–Manin Operator. SIAM Journal on Computing, 36(6), 1777–1806. https://doi.org/10.1137/s0097539704443793
fn main() {
    input!(n: i64);
    if n >= MOD {
        println!("0");
    } else {
        println!("{}", fact(n));
    }
}

まとめ

  • 400行も書いたので流石に疲れた
  • だれか可変modのジャッジのリンクをください

diverta 2019 Programming Contest F - Edge Ordering

5時間45分かけてようやく解けた。本番で解きたかった。
F - Edge Ordering

問題

N頂点M辺の無向グラフがあり、辺には0からM-1までの番号が付いている。0番からN-2番までの辺は全域木をなすことが保証されている。
これらに1からMまでの整数の重みを付けた重み付き無向グラフを考える。ある重み付けについて、それの最小全域木が0番からN-2番までの辺でできた木であるような場合、その重み付けを妥当な重み付けと呼ぶ。妥当な重み付けすべてについて、そのときの最小全域木の重みの総和を求めよ。

解法

Tを、0番からN-2番までの辺からなる全域木とする。
まず、妥当な重み付けがどのような性質を満たすべきかを考える。N-1番からM-1番までの各辺eについて、以下が成り立つことが必要十分であることがわかる:

  • e = (u, v)としたとき、Tの上のu-v間のパスをe_1, ..., e_kとしたとき、e_1, ..., e_kの重みは全てeの重み未満である。

ここで、本来の問題を簡略化した問題を考える。「妥当な重み付けは何通りあるだろうか?」
これには F - Leftmost Ball と同じテクニックが使える。各部分集合 S \subseteq \{0, \ldots, N - 2\}に対して、以下の値がわかればよい:

  • Sのどれかの辺よりも大きい重みがついていなければならない(どれかの辺に依存する)辺の集合を U (\subseteq \{N - 1, \ldots, M - 1\})とする。S + Uを「トポロジカルソート」(上の大小関係を満たすように1から|S|+|U|までの整数で重み付けすること)する方法の個数の、全ての並び替えの個数(|S|+|U|)!に対する割合はいくらか?

これは以下のようなDPでできる:

  • DP1[{}] = 1
  • DP1[  S \cup \{e\}] += DP1[S] / ints[  S \cup \{e\} ] (ここで、 ints[S] := (Sの辺とSのどれかの辺に依存する辺の個数) = (上の|S| + |U|))

DP1[{}] = 1は自明である(空を並べる方法は1通りであり、これはトポロジカルソート)。Sに新しい辺eを追加するときを考える。Sの辺に依存せずeに依存するような辺はints[  S \cup \{e\} ] - 1 - ints[S]本ある。それらをもともとあるints[S]本の辺の中に織り交ぜ、それらの先頭にeをつける方法の割合は1 / (全体の長さ)である。

以上で妥当な重み付けの総数のM!に対する割合が、したがって妥当な重み付けの総数がわかった。本来の問題に戻る。
妥当な重み付けについて、0番からN-2番までの辺についた重みの合計の総和が知りたいのだった。これを求めるために、以下の部分問題を解く必要がある。

  • 自分を含めたx人の人が横一列に並んでおり、自分の左から数えた位置の期待値はkである(1 <= k <= x)。この列にy人がランダムに入る。ランダムに入った後の自分の左から数えた位置の期待値を求めよ。

これを考える。まずkが固定値である場合(自分が必ずk番目にいる場合)を考える。y人それぞれについて、自分の左に入る確率はk / (x + 1)である。よって、自分の位置の増分の期待値はy * k / (x + 1)であり、よって位置の期待値はk * (x + y + 1) / (x + 1)である。
これはkの一次式なので、kが固定値でない場合にも成り立つ。

元の問題に戻る。簡略化した問題のDPを少し変更することで元の問題を解きたい。DP2[S] := (上のS + Uを並べた際の、Sの辺の順位の和 (妥当でない場合は0) の期待値)と置くと、以下のようなDPでできる:

  • DP2[{}] = 0
  • DP2[  S \cup \{e\}] += (DP2[S] * ints[  S \cup \{e\} ] / (ints[S] + 1) + DP1[S] * (|S| + 1)) / ints[  S \cup \{e\} ]

DP2[{}] = 0は自明である(空列において、順位の和は0)。Sに辺eを追加するにあたり、もともとあるints[S]本の前に必ず追加されるのはeの1本で、残りのints[  S \cup \{e\} ] - 1 - ints[S]本はランダムに織り交ぜられる。よって、S内の辺fに対して、もともとの順位の期待値がuだったとすると、新しい順位の期待値は1 + f * ints[  S \cup \{e\} ] / (ints[S] + 1)である。S内に辺は|S|本あるので、期待値は|S| + (もともとの期待値) * ints[  S \cup \{e\} ] / (ints[S] + 1)になる。

残りはintsの計算である。各辺eについて以下の操作ができればよい:

  • e = (u, v)とする。Tの上のu-v間のパスをe_1, ..., e_kとしたとき、S = {e_1, ..., e_k}として T \cap S \neq \emptysetなるすべてのTに対してints[T] += 1をする。

このような操作は高速ゼータ変換で行うことができる。

以上でこの問題を解くことができた。計算量はO(N * 2^N)である。

登場する典型

  • ゼータ変換
    • intsの計算に必要
    • ある部分集合 S \subseteq \{0, \ldots, k - 1\}について、「 T \supseteq SなるすべてのTに対してdp[T] += 1; 」という操作を実行したくなったら考えるべき
  • 確率
    • それっぽい方法はたくさんあるが、大体は誤り
    • 確率を扱うときには、どこが独立か、どこの独立性を無視して良いかというのを正確に理解する
      • 例えば今回は、妥当な重み付けが0番〜N-2番の辺を軽くする方に偏っているので、妥当な重み付けに限ったとき、例えば(0番の重み) < (N番の重み)となる条件付き確率は1/2よりも大きい
  • 期待値と確率が入り混じったDPの遷移
    • 期待値の変換が線形なら、確率と「確率 * 期待値」をもっておけばうまくいく

実装上の注意点

  • 解説では場合の数を求めているが、この解法では確率と期待値を求めている。どちらも等価だが、確率の方はうまく考えないと混乱するので注意。

提出: #5391914 (Rust)

// ModInt省略

fn dfs(r: usize, v: usize, par: usize,
       u: usize, used: &mut [Vec<usize>], g: &[Vec<(usize, usize)>]) {
    used[r][v] = u;
    for &(w, idx) in &g[v] {
        if par == w { continue; }
        dfs(r, w, v, u | 1 << idx, used, g);
    }
}
 
// Depends on ModInt.rs
fn fact_init(w: usize) -> (Vec<ModInt>, Vec<ModInt>) {
    let mut fac = vec![ModInt::new(1); w];
    let mut invfac = vec![0.into(); w];
    for i in 1 .. w {
        fac[i] = fac[i - 1] * i as i64;
    }
    invfac[w - 1] = fac[w - 1].inv();
    for i in (0 .. w - 1).rev() {
        invfac[i] = invfac[i + 1] * (i as i64 + 1);
    }
    (fac, invfac)
}
 
fn construct_ints(n: usize, ab: &[(usize, usize)]) -> Vec<usize> {
    let m = ab.len();
    let mut g = vec![vec![]; n];
    for i in 0..n - 1 {
        let (a, b) = ab[i];
        g[a].push((b, i));
        g[b].push((a, i));
    }
    let mut used = vec![vec![0; n]; n];
    for i in 0..n {
        dfs(i, i, n, 0, &mut used, &g);
    }
    // cnt, ints
    let mut cnt = vec![0usize; 1 << (n - 1)];
    let mut ints = vec![0usize; 1 << (n - 1)];
    for i in 0..m {
        let (a, b) = ab[i];
        cnt[used[a][b]] += 1;
    }
    for i in 0..n - 1 {
        for bits in 0..1 << (n - 1) {
            if (bits & 1 << i) == 0 {
                cnt[bits | 1 << i] += cnt[bits];
            }
        }
    }
    for bits in 0..1 << (n - 1) {
        ints[bits] = m - cnt[(1 << (n - 1)) - 1 - bits];
    }
    ints
}
 
fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize, m: usize,
        ab: [(usize1, usize1); m],
    }
    let (fac, invfac) = fact_init(m + 1);
    let mut inv = vec![ModInt::new(0); m + 1];
    for i in 1..m + 1 {
        inv[i] = fac[i - 1] * invfac[i];
    }
    let ints = construct_ints(n, &ab);
    let mut dp1 = vec![ModInt::new(0); 1 << (n - 1)];
    let mut dp2 = vec![ModInt::new(0); 1 << (n - 1)];
    dp1[0] = ModInt::new(1);
    for bits in 0..1 << (n - 1) {
        for j in 0..n - 1 {
            if (bits & 1 << j) == 0 {
                let x = ints[bits] as i64;
                let y = ints[bits ^ 1 << j] as i64 - x - 1;
 
                let intercalate = inv[(x + y + 1) as usize];
                dp1[bits ^ 1 << j] += dp1[bits] * intercalate;
                let ratio = ModInt::new(x + y + 1) * inv[x as usize + 1];
                dp2[bits ^ 1 << j] += intercalate *
                    (ratio * dp2[bits] + dp1[bits] * (bits.count_ones() as i64 + 1));
            }
        }
    }
    puts!("{}\n", fac[m] * dp2[(1 << (n - 1)) - 1]);
}

まとめ

最後のDP2を計算するパートを思いつくのに一番時間がかかった。解けた後意気揚々と解説を見たらその一番時間がかかるパートが「ほぼ先ほどと同様のDP」で済まされていてキレかけた(は?)。多分諦めて解説を見ていたらキレていたと思う。

Codeforces Global Round 2 H. Triple

本番Eまで解いてF, G, Hのどれを解くか、という状態になったときに一番可能性の高かったこの問題に粘着した。結局時間内には解けなかったが、悔しかったので全力で頑張ったら解けた。
https://codeforces.com/contest/1119/problem/H

問題

整数kが与えられる。数列がn個ある。どの数列も長さ(x + y + z)で、要素は0から2^k-1までの整数である。i番目の数列s[i]はa[i]をx個、b[i]をy個、c[i]をz個含む。
t = 0, 1, ..., 2^k - 1に対して、以下の個数を数え上げて998244353で割った余りを求めよ:

  • 長さnの整数列(u[0], ..., u[n - 1])であって、s[0][u[0]] xor s[1][u[1]] xor ... xor s[n - 1][u[n - 1]] = tを満たすもの。

制約

  • 1 <= n <= 10^5
  • 1 <= k <= 17
  • 0 <= a[i], b[i], c[i] <= 2^k - 1
  • 0 <= x, y, z <= 10^9

解法

各列s[i]について長さ2^kの頻度表を作り、それをv[i]と呼ぶ。v[i]のアダマール変換をh[i]とすれば、求めるものはh[i]の要素ごとの積の逆アダマール変換である。
これを愚直にやるとO(n * k * 2^k)時間かかってしまうので高速化する。一般性を失わず、a[i] = 0としてよい。(b[i]とc[i]にa[i]をxorしておき、最終的な結果にa[i]たちのxorをxorする。) こうすると、アダマール変換後の数列h[i]に登場する数はx ± y ± zという形の4通りのみとなる。それぞれがいつ起こるかは以下の通り:

  • h[i][j] = x - y - z: popcount(b[i] & j)が奇数、popcount(c[i]&j)が奇数
  • h[i][j] = x + y - z: popcount(b[i] & j)が偶数、popcount(c[i]&j)が奇数
  • h[i][j] = x - y + z: popcount(b[i] & j)が奇数、popcount(c[i]&j)が偶数
  • h[i][j] = x + y + z: popcount(b[i] & j)が偶数、popcount(c[i]&j)が偶数

各jに対して、これらが現れる回数を数え上げればよい。
実はこれもアダマール変換を用いることでできる。例えば長さ2^kの配列pを、p[j] = (popcount(b[i] & j)もpopcount(c[i]&j)も奇数 ? 1 : 0)で定めると、配列pは以下によってつくられる長さ2^kの配列qのアダマール変換として得られる:

  • qを0で初期化, q[0] += 1/4, q[b[i] ^ c[i]] += 1/4, q[b[i]] -= 1/4, q[c[i]] -= 1/4

これを利用すれば、毎回配列のたかだか16箇所に足し算をし、最後にまとめてアダマール変換を行うことで、x ± y ± zに乗せるべき指数がわかる。

計算量はO(2^k * k + n)である。

登場する典型

実装上の注意点

  • 手数が多い上に、アダマール変換を使う都合上所々で係数の調整が必要なので、注意して実装する。
  • (x-y-z) % MODみたいな演算をバグらせない
    • MOD < 10^9なので2 * MODを足すだけでは不十分、3*MODを足す (1敗)

提出: #52511941 (Rust)

// ModInt省略

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize, k: usize,
        x: i64, y: i64, z: i64,
        abc: [(usize, usize, usize); n],
    }
    let inv2 = ModInt::new(2).inv();
    let mut base = 0;
    // bias: 4 *
    let mut tbl = vec![vec![0i64; 1 << k]; 4];
    for &(a, b, c) in &abc {
        base ^= a;
        let b = a ^ b;
        let c = a ^ c;
        tbl[0][0] += 4;
        tbl[1][0] += 2;
        tbl[1][b] -= 2;
        tbl[2][0] += 2;
        tbl[2][c] -= 2;
        tbl[3][0] += 1;
        tbl[3][b ^ c] += 1;
        tbl[3][b] -= 1;
        tbl[3][c] -= 1;
    }
    for bits in 0..1 << k {
        for i in 0..2 {
            for j in 0..4 {
                if (j & 1 << i) != 0 {
                    tbl[j ^ 1 << i][bits] -= tbl[j][bits];
                }
            }
        }
    }
    // Hadamard
    for c in 0..4 {
        for i in 0..k {
            for bits in 0..1 << k {
                if (bits & 1 << i) == 0 {
                    let x = tbl[c][bits];
                    let y = tbl[c][bits | 1 << i];
                    tbl[c][bits] = x + y;
                    tbl[c][bits | 1 << i] = x - y;
                }
            }
        }
        for bits in 0..1 << k {
            tbl[c][bits] >>= 2;
        }
    }
    let mut g = vec![ModInt::new(0); 4];
    for kind in 0..4 {
        let mut t = x;
        t += if (kind & 1) == 0 { y } else { -y };
        t += if (kind & 2) == 0 { z } else { -z };
        g[kind] = ModInt::new(t + 3 * MOD);
    }
    let mut prod = vec![ModInt::new(1); 1 << k];
    for bits in 0..1 << k {
        for c in 0..4 {
            prod[bits] *= g[c].pow(tbl[c][bits]);
        }
    }
    // Hadamard
    for i in 0..k {
        for bits in 0..1 << k {
            if (bits & 1 << i) == 0 {
                let x = prod[bits];
                let y = prod[bits | 1 << i];
                prod[bits] = x + y;
                prod[bits | 1 << i] = x - y;
            }
        }
    }
    let fac = inv2.pow(k as i64);
    for bits in 0..1 << k {
        prod[bits] *= fac;
    }
    for bits in 0..1 << k {
        puts!("{}{}", prod[bits ^ base], if bits + 1 == (1 << k) { "\n" } else { " " });
    }
}

まとめ

p[j] = (popcount(b[i] & j)もpopcount(c[i]&j)も奇数 ? 1 : 0)の(逆)アダマール変換が4箇所以下がhotな配列という性質が見えにくかった。
解けたときはドーパミンが多量に分泌されて気持ちよくなった。しかし時間内に解けないのはやはり問題なので、次は時間内に解きたい。

Codeforces Round #520 (Div. 2) F. Upgrading Cities

この手の実装が重い問題を解くのにかなりの困難を感じる。
https://codeforces.com/contest/1062/problem/F

問題

サイクルのないn頂点m辺の有向グラフ(DAG)が与えられる。以下の条件を満たす頂点vを重要な頂点と呼ぶ:

  • 任意の頂点wに対して、vからwへの道かwからvへの道が存在する。

また以下の条件を満たす頂点vを準重要な頂点と呼ぶ:

  • 与えられたグラフからv以外のある頂点一つを取り除いて、vが重要な頂点になるようにできる。

与えられたグラフの重要な頂点と準重要な頂点の個数の合計を求めよ。

制約

  • 2 <= n <= 300000
  • 1 <= m <= 300000
  • 与えられるグラフはDAG

解法

一般性を失わず、与えられたグラフがトポロジカルソートされている、つまり任意の辺(u, v)に対してu < vが成立していることを仮定してよい。
頂点iについて明らかなこととして、頂点iへ到達可能なのはiよりも小さい番号の頂点に限られ、頂点iから到達可能なのはiよりも大きい番号の頂点に限られる。これらは対称的なので、片方だけ考えればよい。つまり、各iについて、i以下の番号の頂点だけ見たときの、(頂点iへ到達不可能な頂点があるかどうか, 適当な1点を取り除いた時に頂点iへ到達不可能な頂点がないようにできるかどうか)がわかれば全てが解決する。

これを求めるにはどうしたら良いだろうか? まず、頂点iへ到達不可能な頂点があるかどうかについては、i以下の番号の頂点だけ見たとき出次数0の頂点がi以外にあるかどうかを判定すれば良い。これはインクリメンタルに出次数を管理していくことで、全てのiについて合計O(n+ m)時間で計算できる。
適当な1点を取り除いた時に頂点iへ到達不可能な頂点がないようにできるかどうかは少し難しい。まず、i以下の番号の頂点だけ見たときに、出次数0の頂点が2個以上であれば不可能であり、0個なら可能である。出次数0の頂点が1個の場合のみが問題だが、その頂点をxとした時、xへ向かう辺(w, x)があるようなすべての頂点wについて、xがなくなった時の出次数が0以外、つまり元々の出次数が2以上であればよい。これもデータ構造をうまく設計すればでき、各頂点vについて「vへの辺がある頂点の中で次数1のものが何個あるか」を管理すればよい。これも全てのiについて合計O(n + m)時間で行える。

同じことを辺を逆向きに張ったグラフに対しても行えば、合計O(n + m)時間で各頂点についての判定が行える。

登場する典型

  • トポロジカルソート
  • 列を真ん中で区切って時系列データ2個に分ける
    • iへ到達可能なのはiより小さい頂点だけ、iから到達可能なのはiより大きい頂点だけなので、綺麗に分割できる
    • 分割した後の両方とも、小さい方から順番に作っていくインクリメンタルなアルゴリズムで計算できる

実装上の注意点

  • 実装が重いので注意
    • データ構造を管理するところで非常に間違えやすいので、デバッグプリントなどを使い正しいことを確かめる
  • 逆辺を張ったグラフについても同じことをやるので、うまく共通化する
  • トポロジカル順序を考える必要があると大変なので、トポロジカル順序を使ってグラフそのものを書き換える
    • 任意の辺(u, v)に対してu < vとなるようにしておくとかなり楽
  • HashSetよりBTreeSetの方が速い
    • おそらく使い方の問題で、要素の追加、削除、1点集合の時に唯一の要素を取り出す、の3種類の操作だけが行われ、HashMapだと3番目の操作が遅い

提出: #51977370 (Rust, HashSet版)
#51982893 (Rust, BTreeSet版)

fn calc(g: &[Vec<usize>], revg: &[Vec<usize>], cnt: &mut [usize]) {
    let n = revg.len();
    let mut outdeg = vec![0; n];
    let mut zeroback = vec![0; n]; // zeroback[v] = #{w | (v -> u) in g, deg[v] = 1}
    let mut zero = 0;
    let mut zs = BTreeSet::new();
    for v in 0..n {
        for &w in &revg[v] {
            if outdeg[w] == 0 {
                zero -= 1;
                zs.remove(&w);
                zeroback[v] += 1;
            }
            if outdeg[w] == 1 {
                for &k in &g[w] {
                    if k < v {
                        zeroback[k] -= 1;
                    }
                }
            }
            outdeg[w] += 1;
        }
        cnt[v] += zero;
        if zero == 1 {
            let &x = zs.iter().next().unwrap();
            cnt[v] += zeroback[x];
        }
        zero += 1;
        zs.insert(v);
    }
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        uv: [(usize1, usize1)],
    }
    let mut g = vec![vec![]; n];
    let mut revg = vec![vec![]; n];
    for &(u, v) in &uv {
        g[u].push(v);
        revg[v].push(u);
    }
    // find a topological order
    let mut ord = vec![];
    {
        let mut indeg = vec![0; n];
        let mut que = VecDeque::new();
        for i in 0..n {
            indeg[i] = revg[i].len();
            if revg[i].is_empty() {
                que.push_back(i);
            }
        }
        while let Some(v) = que.pop_front() {
            ord.push(v);
            for &w in &g[v] {
                indeg[w] -= 1;
                if indeg[w] == 0 {
                    que.push_front(w);
                }
            }
        }
    }
    assert_eq!(ord.len(), n);
    // Replace g and revg
    {
        let mut invord = vec![0; n];
        for i in 0..n {
            invord[ord[i]] = i;
        }
        for i in 0..n {
            g[i].clear();
            revg[i].clear();
        }
        for &(u, v) in &uv {
            let a = invord[u];
            let b = invord[v];
            g[a].push(b);
            revg[b].push(a);
        }
    }
    let mut cnt = vec![0; n];
    calc(&g, &revg, &mut cnt);
    cnt.reverse();
    for i in 0..n {
        for w in g[i].iter_mut() {
            *w = n - 1 - *w;
        }
        for w in revg[i].iter_mut() {
            *w = n - 1 - *w;
        }
    }
    g.reverse();
    revg.reverse();
    calc(&revg, &g, &mut cnt);
    cnt.reverse();
    let mut ans = 0;
    for i in 0..n {
        if cnt[i] <= 1 {
            ans += 1;
        }
    }
    puts!("{}\n", ans);
}

まとめ

ARC092 F - Two Faced Edges - koba-e964の日記にも通じるものがある気がする。今回は自力で解けたので嬉しかった。ただダラダラしてしまい実装に1.5時間近く掛かったのは反省。

yukicoder No.803 Very Limited Xor Subset

この手の問題は最近割と見る印象がある。
No.803 Very Limited Xor Subset - yukicoder

問題

長さNの整数列Aがある。以下の条件を満たす部分列Sの個数をmod 10^9 + 7で求めよ:

  • Sのxor和はX
  • Sは以下のM個の条件(i = 1, ..., M)を満たす
    • L[i]番目からR[i]番目の要素の中でSに含まれるものの個数を2で割った余りはtype[i]である

制約

  • 1 <= N <= 300
  • 0 <= M <= min(300, N(N - 1) / 2)
  • 0 <= X, A[i] <= 10^9
  • 1 <= L[i] <= R[i] <= N
  • 入力は全て整数

解法

A[i]やXを長さ30の0/1ベクトルとみなす。選ぶ個数についての制約は次のようにエンコードできる:

  • j個目の制約用に、A[i]やXに列を追加する(30+j列目になる)。A[i]については、iがL[j]からR[j]に含まれていたらその列は1にする。Xについては、type[j]にする。

こうした場合、元の問題で全ての制約を満たすことと、この問題でAの部分列Sの要素ごとのxorがXになっていることは同値。そのようなSの個数を数え上げればよい。

線形代数を使うと、求めるSの個数はA[i]たちで生成できる空間にXが要素として含まれていたら|ker A| = 2^{dim (ker A)}、含まれていなければ0であることがわかる。この手の計算はGauss-Jordanの消去法ででき、計算量はO(N * M * min(N, M))である。

登場する典型

  • 使用するかしないかをGF(2)の要素と思って線形代数をする
    • Gauss-Jordanの消去法
  • rank A + dim (ker A) = dim cod A
    • 特定の0/1ベクトルをxor和として実現する方法を数えたい場合にかなりよく見るテク
  • 制約をエンコードするために列を追加

実装上の注意点

  • Gauss-Jordanの消去法はそこそこ間違えやすいので気をつける
  • 数え上げだが実質的にはdim kerを求める問題なので、サンプルがあっているからといって過信しない

提出: #325170 (Rust)

// ModInt省略

// returns (rank(a), dim ker(a))
fn gauss_elim_gf2(a: &mut [Vec<i32>], b: &mut [i32]) -> Option<(usize, usize)> {
    let n = a.len();
    let m = b.len();
    let mut row = 0;
    for col in 0..m {
        if row >= n { break; }
        let mut row2 = None;
        for i in row..n {
            if a[i][col] == 1 {
                row2 = Some(i);
                break;
            }
        }
        if row2.is_none() { continue; }
        let row2 = row2.unwrap();
        a.swap(row, row2);
        for k in row + 1..n {
            if a[k][col] == 1 {
                for j in 0..m {
                    a[k][j] ^= a[row][j];
                }
            }
        }
        if b[col] == 1 {
            for j in 0..m {
                b[j] ^= a[row][j];
            }
        }
        row += 1;
    }
    for i in 0..m {
        if b[i] != 0 { return None; }
    }
    Some((row, n - row))
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        m: usize,
        x: i64,
        a: [i64; n],
        tlr: [(i32, usize1, usize); m],
    }
    let mut mat = vec![vec![0; 30 + m]; n];
    for i in 0..n {
        for j in 0..30 {
            mat[i][j] = if (a[i] & 1 << j) != 0 { 1 } else { 0 };
        }
    }
    let mut b = vec![0; 30 + m];
    for i in 0..30 {
        b[i] = if (x & 1 << i) != 0 { 1 } else { 0 };
    }
    for i in 0..m {
        let (t, l, r) = tlr[i];
        for j in l..r {
            mat[j][30 + i] = 1;
        }
        b[30 + i] = t;
    }
    match gauss_elim_gf2(&mut mat, &mut b) {
        Some((_, ker)) => {
            puts!("{}\n", ModInt::new(2).pow(ker as i64));
        },
        None => {
            puts!("0\n");
        },
    }
}

まとめ

この手の線形代数をやる問題は結構すき (解けるので)

yukicoder No.802 だいたい等差数列

こういうのはかなり得意だと思っている。
No.802 だいたい等差数列 - yukicoder

問題

以下の条件を満たす長さNの整数列Aの個数を数え、10^9+7で割った余りを求めよ。

  • 1 <= A[1] <= ... <= A[N] <= M
  • i = 1, ..., N - 1に対して D1 <= A[i + 1] - A[i] <= D2

制約

  • 2 <= N <= 3 * 10^5
  • 1 <= M <= 10^6
  • 0 <= D1 <= D2 <= M

解法

まず愚直DPを考える。DP[i][j] := (Aの第i項まで見たとき、A[i]=jとなるような数列の個数)としたとき、以下のような遷移を持つDPを行えばよい。

  • init: DP[1][j] = 1
  • step: DP[i + 1][j + k] += DP[i][j] (D1 <= k <= D2)

これのDP[N][1] + ... DP[N][M]が答えである。これを実行することは計算量が多すぎて不可能なので、高速化を行う。
DP配列を多項式と考えることにして、f(i) = \sum DP[i][j] * x^jと置く。(最終的に興味があるのはx^M以下の項なので、x^{M+1}以上の項は無視する。また多項式級数(項が無限個続くもの)を断りなく同一視する。)
最初はf(1) = x + x^2 + ... + x^M = x/(1 - x) + O(x^{M + 1})である。
またstepを多項式の言葉に直すとf(i + 1) = f(i) * (x^{D_1} + ... + x^{D_2}) = f(i) * (x^{D_1} - x^{D_2 + 1}) / (1 - x)である。
また、多項式A + Bx + ... + Cx^Mに1 / (1 - x)を掛けたもの(のマクローリン展開)のx^Mの係数はA + B + ... + Cである。
よってこれらを組み合わせると、x * (x^{D_1} - x^{D_2 + 1})^{N - 1} / (1 - x)^{N + 1} (をマクローリン展開したもの)におけるx^Mの係数が答えであることがわかる。
これはどう計算すべきだろうか? まず、(x^{D_1} - x^{D_2 + 1})^{N - 1} = \sum_{i=0}^{N - 1} C(N - 1, i) * (-1)^i * x^{D_1 * (N - 1 - i) + (D_2 + 1) * i)}なので、「ある整数kに対してx^k / (1 - x)^{N + 1}におけるx^Mの係数」を高速に計算できればよいことがわかる(このようなクエリはN回実行される)。
これは1 / (1 - x)^{N + 1}におけるx^{M - k}の係数に他ならないので、M - k < 0であれば0で、M - k >= 0であればC(M - k + N, N)である。

登場する典型

  • 母関数
    • 愚直DPの高速化

実装上の注意点

  • 二項係数の扱いに気をつける
  • 母関数を使う時は式変形が複雑になりがちなので、式変形が間違っていないかどうかサンプルなどを使って確かめる

提出: #324980 (Rust)

fn fact_init(w: usize) -> (Vec<ModInt>, Vec<ModInt>) {
    let mut fac = vec![ModInt::new(1); w];
    let mut invfac = vec![0.into(); w];
    for i in 1 .. w {
        fac[i] = fac[i - 1] * i as i64;
    }
    invfac[w - 1] = fac[w - 1].inv();
    for i in (0 .. w - 1).rev() {
        invfac[i] = invfac[i + 1] * (i as i64 + 1);
    }
    (fac, invfac)
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        m: usize,
        d1: usize,
        d2: usize,
    }
    let n = n - 1;
    let (fac, invfac) = fact_init(n + m + 2);
    let mut tot = ModInt::new(0);
    for i in 0..n + 1 {
        let deg = d1 * (n - i) + (d2 + 1) * i + 1;
        if deg > m { continue; }
        let rem = m - deg;
        let tmp = fac[rem + n + 1] * invfac[n + 1] * invfac[rem]
            * fac[n] * invfac[n - i] * invfac[i];
        if i % 2 == 0 {
            tot += tmp;
        } else {
            tot -= tmp;
        }
    }
    puts!("{}\n", tot);
}

まとめ

母関数周りも今後定跡になりそうな気がする。

早稲田大学プログラミングコンテスト2019 D - Choose Your Characters

さすがにこれが解けないのは反省。
D - Choose Your Characters

問題

N種類のキャラクターがいるゲームで対戦を行う。N種類のキャラクターには1からNまでの番号がついている。またそれらの間には相性という概念が定まっており、「キャラクターaはキャラクターbに対して有利(キャラクターbはキャラクターaに対して不利)」という関係がM個存在する。
このゲームでA君とB君が対戦する。A君にどうしても勝ちたいB君は、以下のような条件で特訓を行うことにした。

  • 特訓するキャラクターは[L, R] (閉区間)の範囲の番号のキャラクターに集中して行う。少なければ少ないほどよい。
  • A君がどのキャラクターaを選んでも、[L, R]内のキャラクターbであって、a != bでbはaに対して不利ではないようなbが存在する
    • 少なくとも五分五分以上にはなるようなbが存在する

条件を満たす閉区間[L, R]は存在するか? 存在する場合は長さがもっとも短いものを、複数存在する場合はLが最小のものを出力せよ。存在しない場合は-1を出力せよ。

制約

  • 2 <= N <= 100000
  • 1 <= M <= 200000
  • 同じ条件は1回まで出現する
  • 正反対の条件は出現しない

解法

二分探索を使うと、区間長R - L + 1を固定した時に、条件を満たすようなLが存在するかどうかをO(N)時間程度で判定できればよいことがわかる。
これは可能で、以下のようにすれば良い。

まず、各キャラクターbについて、B君がbを選んだ時にA君に選んでほしくないキャラクター(bが不利になるキャラクターとb自身)を管理する。
以降B君が選ぶ区間[L, R]を伸び縮みさせることを考える。[L, R]の中のキャラクターに対するA君に選んで欲しくないキャラクターが、何回数えられるかを管理する。明らかに、選んで欲しくないキャラクターとしてR - L + 1回登場するキャラクターが存在する場合、今見ている範囲[L, R]ではそのキャラクターには対応できない。
長さlenで条件を満たす区間が存在するかどうか確かめる場合には、[0, len - 1] -> [0, len] -> [1, len] -> [1, len + 1] -> ... のように、右端と左端を交互に右にずらしていけばよい。あるステップで区間に入ったり区間から除外される頂点をvとすれば、そのステップの計算量はO(vを選んだ時にA君に選んでほしくないキャラクターの個数)である。その合計はO(M)なので、長さlenでOKかどうかの判定もO(M)時間である。

2分探索なので全体の計算量はO(M log N)。

登場する典型

  • 尺取り法
    • セグメント木を使いたいが計算速度が間に合わないときは、尺取り法やスライド最小値でなんとかできないかを考える
  • 頻度表
    • setを使いたくなったら頻度表でなんとかできないかを考える

実装上の注意点

  • 尺取りなので気をつける
    • 2分探索時とLを求める時の2回行うので、関数にした方がよい
  • 頻度表なのでデータの整合性に気をつける
    • 変換用の関数を作るとよい

提出: #4543465 (Rust)

fn op(len: usize, a: &mut [usize], count: &mut usize, idx: usize, newval: usize) {
    if a[idx] == len {
        *count -= 1;
    }
    a[idx] = newval;
    if newval == len {
        *count += 1;
    }
}

fn sweep(g: &[Vec<usize>], len: usize) -> Option<usize> {
    let n = g.len();
    // sweep
    let mut freq = vec![0; n];
    let mut cnt = 0;
    for i in 0..len {
        for &v in &g[i] { 
            let val = freq[v] + 1;
            op(len, &mut freq, &mut cnt, v, val);
        }
    }
    for i in 0..n - len + 1 {
        if cnt == 0 {
            return Some(i);
        }
        // progress
        if i < n - len {
            for &v in &g[i + len] {
                let val = freq[v] + 1;
                    op(len, &mut freq, &mut cnt, v, val);
                }
            for &v in &g[i] {
                let val = freq[v] - 1;
                op(len, &mut freq, &mut cnt, v, val);
            }
        }
    }
    None
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        m: usize,
        ab: [(usize1, usize1); m],
    }
    let mut g = vec![vec![]; n];
    for &(a, b) in &ab {
        g[b].push(a);
    }
    for i in 0..n {
        g[i].push(i);
    }
    let mut pass = n + 1;
    let mut fail = 1;
    while pass - fail > 1 {
        let mid = (pass + fail) / 2;
        if sweep(&g, mid).is_some() {
            pass = mid;
        } else {
            fail = mid;
        }
    }
    // len = pass
    if pass > n {
        puts!("-1\n");
        return;
    }
    let l = sweep(&g, pass).unwrap();
    puts!("{} {}\n", l + 1, l + pass);
}

まとめ

本番は「セグ木に平衡2分探索木を乗せて、毎回O(N log^2 N)時間かけて判定すればいいから全体ではO(N log^3 N)!w」みたいなことを考えて途方にくれていたが、セグ木も平衡2分探索木も自然により良い方法が存在したので、考えるべきだった。