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
投稿日時:
最終更新:
