https://atcoder.jp/contests/abc456/tasks/abc456_g
久しぶりに解説を見てみたら、母関数を使う方法が書いてありました。
Sをxで分割して、それぞれの領域の長さをとして、
日のうち連続する休日が
日までの場合の数を
とすると、求める場合の数は、
となります。なので、このを求めればよいです。
休日をo、勤務日をxで表すと、oは個まで続いていいので、x、xo、…、xo…o(oが
個)を組み合わせて結合すればいいです。ただしこれだと必ずxから始まるので
の長さを
までの長さで分割することにします。例えば長さ3を1以下で分割するなら、
1 1 1 1 => x x x x => xxx 2 1 1 => xo x x => oxx 1 2 1 => x xo x => xox 1 1 2 => x x xo => xxo 2 2 => xo xo => oxo
の5通りになります。
一般化のために個を
個以内で分割することを考えて、これを母関数にすると、
となって、このの係数が求める場合の数となります。
個に分割して、
番目の長さが
です。
ここまでは考えていたのですが、内側の和がで終わるのが怖すぎて、ここでストップしてしまいました。解説ではここからさらに計算しています。
こうしてみるとが大きいと速く計算できそうな気がします。だから、
でくくると、
ここでを考えます。その前に
を考えると、これは
なので、
の係数は
種類の中から
個を選ぶ重複組み合わせになります。
なら
を
に置き換えるだけなので、
だから個別の係数はすぐに計算できます。なので、の
の係数は、
で計算できます。
// Count Holidays #![allow(non_snake_case)] //////////////////// library //////////////////// fn read<T: std::str::FromStr>() -> T { let mut line = String::new(); std::io::stdin().read_line(&mut line).ok(); line.trim().parse().ok().unwrap() } // ax = by + 1 (a, b > 0) fn linear_diophantine(a: i64, b: i64) -> Option<(i64, i64)> { if a == 1 { return Some((1, 0)) } let q = b / a; let r = b % a; if r == 0 { return None } let (x1, y1) = linear_diophantine(r, a)?; Some((-q * x1 - y1, -x1)) } fn inverse(a: i64, d: i64) -> i64 { let (x, _y) = linear_diophantine(a, d).unwrap(); if x >= 0 { x % d } else { x % d + d } } fn pow(n: i64, e: u32, d: i64) -> i64 { if e == 0 { 1 } else if e == 1 { n } else if e % 2 == 1 { pow(n, e-1, d) * n % d } else { let m = pow(n, e/2, d); m * m % d } } //////////////////// HCalculator //////////////////// const D: i64 = 998244353; struct HCalculator { facs: Vec<i64>, inv_facs: Vec<i64> } impl HCalculator { fn calc(&self, n: i64, m: i64) -> i64 { self.facs[(n+m-1) as usize] * self.inv_facs[m as usize] % D * self.inv_facs[(n-1) as usize] % D } fn build(N: usize) -> HCalculator { let mut facs: Vec<i64> = vec![1; N+1]; let mut inv_facs: Vec<i64> = vec![1; N+1]; for n in 2..N+1 { facs[n] = facs[n-1] * (n as i64) % D; inv_facs[n] = inv_facs[n-1] * inverse(n as i64, D) % D } HCalculator { facs, inv_facs } } } //////////////////// process //////////////////// fn read_input() -> String { let _N: i64 = read(); let S: String = read(); S } fn divide(S: String) -> Vec<i64> { let mut ns: Vec<i64> = vec![]; let mut n: i64 = 0; for c in S.chars() { if c == '.' { n += 1 } else { if n > 0 { ns.push(n) } n = 0 } } if n > 0 { ns.push(n) } ns } use std::collections::BTreeMap; fn frequency(v: &Vec<i64>) -> Vec<(i64, u32)> { let mut m: BTreeMap<i64, u32> = BTreeMap::new(); for &n in v.iter() { let e = m.entry(n).or_insert(0); *e += 1 } m.into_iter().rev().collect::<Vec<_>>() } // 2^e % D fn pow2(e: i64) -> i64 { if e == 0 { 1 } else if e == 1 { 2 } else if e % 2 == 1 { pow2(e-1) * 2 % D } else { let n = pow2(e/2); n * n % D } } // nを1~kの整数で分割するときの場合の数 fn num_divs(n: i64, k: i64, hc: &HCalculator) -> i64 { let mut s: i64 = 0; for p in 0..(n-1)/(k+1)+1 { let sign: i64 = if p % 2 == 0 { 1 } else { -1 }; // pのときの最小次数 let q1 = (k+1)*p+1; let m1 = n - q1; s = (s + sign * pow2(m1) * hc.calc(p+1, m1)).rem_euclid(D); let q2 = (k+1)*(p+1); let m2 = n - q2; if m2 >= 0 { s = (s - sign * pow2(m2) * hc.calc(p+1, m2)).rem_euclid(D) } } s } fn F(S: String) { let N = S.len(); let ns = divide(S); let freq = frequency(&ns); let hc = HCalculator::build((N*2) as usize); let mut accs: Vec<(Vec<i64>, u32)> = vec![]; for (n, f) in freq { let acc: Vec<i64> = (0..n+1).map(|k| num_divs(n+1, k+1, &hc)).collect(); accs.push((acc, f)) } // i日連続以下の場合の数 let mut a: Vec<i64> = vec![1; N+1]; let mut full: i64 = 1; for i in 1..N+1 { a[i] = full; for (acc, f) in accs.iter() { if i >= acc.len() { break } a[i] = a[i] * pow(acc[i], *f, D) % D; if i == acc.len() - 1 { // ここで最後、あとは同じ値が続くとする full = full * pow(acc[i], *f, D) % D } } } for i in 1..N+1 { println!("{}", (a[i] - a[i-1]).rem_euclid(D)) } } fn main() { let S = read_input(); F(S) }