Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 276 G問題 Count Sequences

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

次を満たす長さ  N の整数列  A=(A_1, \dots, A_N) の個数を求めよ.

  •  0 \leq A_1 \leq A_2 \leq \dots \leq A_N \leq M
  •  i=1,2, \dots, N-1 に対して,  A_i 3 で割った余りと  A_{i+1} 3 で割った余りが異なる.

制約

  •  2 \leq N \leq 10^7
  •  1 \leq M \leq 10^7

解法 1 (解法の方針建て)

形式的べき級数を用いて考えることにする. 次のようにして整数列  A を構成することを考える.

  • 初項は非負整数を適当に選ぶ
  • それ以降は末項  a に3の倍数ではない整数  m を選び,  (a+m) を末尾に追加する.
  • 長さが  N になった時, 末項が  M 以下ならば成功, そうでなければ失敗

このとき, 条件を満たすような整数列と上の作り方で成功したときの整数列は操作の方法と1対1に対応する.

「初項は非負整数を適当に選ぶ」に対応する形式的べき級数  f(x)

 \displaystyle f(x)=\sum_{n=0}^\infty x^n=\dfrac{1}{1-x}

である.

「それ以降は末項  a に3の倍数ではない整数  m を選び,  (a+m) を末尾に追加する.」に対応するのは形式的べき級数  g(x)

 \displaystyle g(x)=\sum_{\substack{0 \leq n \\ n \not \equiv 0 \pmod{3}}} x^n=\sum_{n=0}^\infty x^n-\sum_{m=0}^\infty x^{3m}=\dfrac{1}{1-x}-\dfrac{1}{1-x^3}=\dfrac{x(1+x)}{1-x^3}

としたとき,  g(x) を掛けることになる.

このとき, 末項が  m である整数列の数は

 \displaystyle \left[x^m \right] (f g^{N-1})

である. これを  m=0,1,2, \dots, M とした総和をとることにより,

 \displaystyle \sum_{m=0}^M \left[x^m \right] (fg^{N-1})=\left[x^M \right] \dfrac{fg^{N-1}}{1-x}=\left[x^M \right] f^2 g^{N-1}

となる.

よって,  f^2 g^{N-1} M 次の係数を求めることを目標にする.

解法 2 (形式的べき級数の処理)

 f^2 g^{N-1}=\dfrac{1}{(1-x)^2} \dfrac{x^{N-1} (1+x)^{N-1}}{(1-x^3)^{N-1}}=x^{N-1} \times \dfrac{(1+x)^{N-1}}{(1-x)^2} \times \dfrac{1}{(1-x^3)^{N-1}}

であるから,

 \begin{align*}
\displaystyle \left[x^M \right] f^2 g^{N-1}
&=\left[x^M \right] \left(x^{N-1} \times \dfrac{(1+x)^{N-1}}{(1-x)^2} \times \dfrac{1}{(1-x^3)^{N-1}}\right) \\
&=\left[x^{M-(N-1)} \right] \left(\dfrac{(1+x)^{N-1}}{(1-x)^2} \times \dfrac{1}{(1-x^3)^{N-1}}\right) \end{align*}

である.  K:=M-(N-1) とする.  K \lt 0 ならば, 解答は  0 である. 以降は  K \geq 0 であるとする.

2つの形式的べき級数  p,q をそれぞれ

 p(x):=\dfrac{(1+x)^{N-1}}{(1-x)^2}, \quad q(x):=\dfrac{1}{(1-x^3)^{N-1}}

とする.

 p(x) について,  h(x):=(1+x)^{N-1}

 \displaystyle h(x)=\sum_{n=0}^{N-1} \dbinom{N-1}{n} x^n

である. そして, "形式的べき級数 (1-x) で割るということ" と, "各係数の累積和に対応する形式的べき級数を生成する" ということが等価になる. よって,  p(x)=h(x)/(1-x)^2 であったから,  h の係数の累積和をとるということを2回連続で行うことにより,  p (N-1) 次までの係数を得ることが出来る.

 q(x) については, 負の二項定理 より,

 \displaystyle q(x)=\sum_{j=0}^\infty \dbinom{j+(N-2)}{N-2} x^{3j}

である.

以上から,  p,q i 次の係数を  p_i,  j 次の多項式 q_j と書くことにすると, 求めるべきは

 \displaystyle \left[x^K \right] (pq)=\sum_{\substack{0 \leq i,j \\ i+3j=K}} p_i q_j=\sum_{\substack{0 \leq i,j \\ i+3j=K}} p_i \dbinom{j+(N-2)}{N-2}

である.

計算量は  i!~(0 \leq i \leq M) O(M) 時間での前計算により,  h を求めるパート, 累積和を求めるパート,  q を求めるパート,  \left[x^K \right ] (pq) を求めるパート全て  O(M) 時間で計算できる.