F - Random Shuffles 解説 by ngtkana


近似を陽に使わない解法のうち、公式解説で紹介されていた最速のものは \(\Theta ( N D \lg K )\) 時間でした。このユーザー解説では \(\Theta ( N D )\) 時間のものをご紹介します。(定数倍はかなり厳しいです。)

さて、「小課題 1-10 (800 点): \(N \le 800 000, D \le 10, K \le 10^9\)」の公式解説の内容の要点は、長さ \(D\) の数列の積

\[ (a * b) (k) = \sum _ { i + j \equiv k \pmod D } a ( i ) b ( j ) \quad (a, b \in \mathbb R ^ D, 0 \le k \lt D ) \]

に関する \(L\) 乗の計算を \(L\) に関する二分累乗により \(\Theta ( D ^ 2 \lg L )\) で行うことでした。ところでこれは巡回畳み込みそのものですから、複素数に拡張して DFT(離散フーリエ変換)をすることで数列の \(L\) 乗が各点 \(L\) 乗に移り、\(\Theta ( D ^ 2 )\) (長さが2冪などとは限らないので FFT を使えませんでした。)で計算できることがわかります。

use itertools::Itertools;
use nalgebra::{Complex, ComplexField};
use proconio::{fastout, input};

#[fastout]
fn main() {
    input! {
        _n: usize,
        d: usize,
        l: usize,
        s: String,
    }

    let mut weight = vec![vec![0; d]; d];
    let mut ans = vec![0.0; d];
    for (chunk_id, chunk) in s
        .chars()
        .map(|c| c.to_digit(10).unwrap() as usize)
        .enumerate()
        .chunks(d)
        .into_iter()
        .enumerate()
    {
        let mut a = vec![0.0; d];
        let p = ((d * (1 + chunk_id)) as f64).recip();
        a[0] = 1.0 - p - p;
        a[1] = p;
        a[d - 1] = p;
        let a = pow(&a, l);
        for (i, x) in chunk {
            weight[x][(i + d - x) % d] += i + 1;
        }
        for (i, wi) in weight.iter().enumerate() {
            ans[i] += wi.iter().zip(&a).map(|(&x, &y)| x as f64 * y).sum::<f64>();
        }
    }
    for &ans in &ans {
        println!("{}", ans);
    }
}

fn pow(a: &[f64], b: usize) -> Vec<f64> {
    let n = a.len();
    let a = a.iter().map(|&x| Complex::from_real(x)).collect_vec();
    let a = (0..n)
        .map(|i| {
            let w = Complex::from_polar(&1.0, &(2.0 * std::f64::consts::PI * i as f64 / n as f64));
            let mut coeff = Complex::from_real(1.0);
            let mut ans = Complex::from_real(0.0);
            for x in a.iter() {
                ans += x * coeff;
                coeff *= w;
            }
            ans.powf(b as f64)
        })
        .collect_vec();
    (0..n)
        .map(|i| {
            let w =
                Complex::from_polar(&1.0, &(2.0 * (-std::f64::consts::PI * i as f64) / n as f64));
            let mut coeff = Complex::from_real(1.0);
            let mut ans = 0.0;
            for x in a.iter() {
                ans += (x * coeff).real();
                coeff *= w;
            }
            ans / n as f64
        })
        .collect_vec()
}

提出 (944 ms): https://atcoder.jp/contests/s8pc-6/submissions/28576965

投稿日時:
最終更新: