包除原理

この問題で数え上げる対象は崩した言い方で言うと

• (\(1\) 番目の庭に \(1\) 個以上の苗を植える) and
  (\(2\) 番目の庭に \(1\) 個以上の苗を植える) and
  (\(3\) 番目の庭に \(1\) 個以上の苗を植える) …

という感じになります。しかしこれを直接数え上げるのはかなり困難で、ラフな動的計画法では \(\mathrm{O}(NABC) = \mathrm{O}(N^4)\) 程度の計算量が掛かってしまいます。
一方、もしも上の 「植える」の部分を「植えない」に変えたらとたんに数えやすくなり、 \(1\) 通りが答えです。
このように、「数え上げの対象が直接計算しようとするとかなり数えづらいが、 and を or に変えたり \(A\) を not \({A}\) に変えたりすると数え上げやすくなる」という場合は包除原理の出番であることが多いです。

包除原理の公式

\[ \left|\bigcap_{i \in S} A_i \right| = \sum_{T \subseteq S} (-1)^{|T|} \left| \bigcap_{i \in T} \overline{A_i} \right|\]

に \(A_i :=\) ( \(i\) 番目の庭に \(1\) 個以上の苗を植える置き方) を代入します。 すると、求めるのは砕けた表現で書くと

\[ \begin{aligned} &(すべての庭に 1 個以上苗が植えられている) \\ &= (少なくとも特定の 0 個の庭に苗が植えられていない) \\ &- (少なくとも特定の 1 個の庭に苗が植えられていない) \\ &+ (少なくとも特定の 2 個の庭に 苗が植えられていない) \\ &- \dots \\ &= \sum_{i=0}^N (-1)^i (少なくとも特定の i 個の庭に苗が植えられていない) \end{aligned} \]

という感じになります。

(少なくとも特定の \(i\) 個の庭に苗が植えられていない) 場合の数を求めると、

• 植えられていない庭の選び方 \(\to\) \(\binom{N}{i}\) 通り
• リンゴの苗の植え方 \(\to\) 残りの \(N-i\) 個の庭に \(0\) or \(1\) 個ずつ植えるかどうか選べる。\(a\) 本植えるとして \(\binom{N-i}{a}\) 通りなのでこれを \(0 \leq a \leq \min(A,N-i)\) について足し合わせたものが通り数になる

という考察により

\[\binom{N}{i} \sum_{a=0}^{\min(A,N-i)} \binom{N-i}{a} \sum_{b=0}^{\min(B,N-i)} \binom{N-i}{b} \sum_{c=0}^{\min(C,N-i)} \binom{N-i}{c}\]

とわかります。以上より答えは

\[ \begin{aligned} &\sum_{i=0}^N (-1)^i \binom{N}{i} \sum_{a=0}^{\min(A,N-i)} \binom{N-i}{a} \sum_{b=0}^{\min(B,N-i)} \binom{N-i}{b} \sum_{a=0}^{\min(C,N-i)} \binom{N-i}{c} \\ &=\sum_{i=0}^N (-1)^{N-i} \binom{N}{i} \sum_{a=0}^{\min(A,i)} \binom{i}{a} \sum_{b=0}^{\min(B,i)} \binom{i}{b} \sum_{c=0}^{\min(C,i)} \binom{i}{c} \end{aligned} \]

になるとわかりました。

上式はそのまま計算しようとすると \(\mathrm{O}(N^2)\) かかり、計算量は大きく改善できたものの残念ながらこのままでは TLE してしまうので更なる工夫が必要です。

グリッドを利用した高速化

困ったようですが、実は上式はグリッドを用いた言い換えを利用することで計算量を \(\mathrm{O}(N)\) に落とすことができます。

グリッドを利用すると、事前に \(1!,2!, \dots,N!\) と \(\frac{1}{1!}, \frac{1}{2!}, \dots, \frac{1}{N!}\) を前計算してあるときに

\[f_M(N) = \sum_{i=0}^{\min(N,M)} \binom{N}{i}\]

から

\[f_M(N + 1) = \sum_{i=0}^{\min(N+1,M)} \binom{N+1}{i}\]

を \(\mathrm{O}(1)\) で求めることができて、\(f_M(0), f_M(1), \dots, f_M(N)\) を \(\mathrm{O}(N)\) で列挙することができます。さっそくこの方法を説明していきましょう。

まず、\(N + 1 \leq M\) のときは単純で、

\[f_M(N + 1) = \sum_{i=0}^{N+1} \binom{N+1}{i} = 2^{N+1}\]

なので

\[f_M(N + 1) = 2 \times f_M(N)\]

です。
問題は \(N + 1 \gt M \iff N \ge M\) ですが、ここでグリッドを持ち出しましょう。

「グリッドの左下を始点として、右あるいは上に \(1\) 歩ずつ移動できるときに、ある地点に移動する場合の数は？」という問題を考えます。以下の図において

• \(N\) 回移動して赤丸がある格子点へ行く歩き方の場合の数

と

• \(N+1\) 回移動して青丸がある格子点へ行く歩き方の場合の数

の間にはどのような関係が成り立つでしょうか？

答えとしては、赤丸がある地点から右か上に \(1\) 回移動すると、一部の例外を除けば青丸がある地点にいけることから、

\[ \begin{aligned} &(青丸がある格子点への行き方) \\ &= 2 \times (赤丸がある格子点の行き方) - (一番下の赤丸がある格子点から右へ行く歩き方) \end{aligned} \]

という式を得ます。この式を \(f_M(N),f_M(N+1)\) および二項係数を用いた表現に置き換えます。\(\binom{a}{b}\) は「グリッドの左下を出発して、右に \(b\) 歩、上に \(a - b\) 歩歩く方法の場合の数」と言い換えられるのを利用すると、

• \(f_M(N), f_M(N + 1)\) はそれぞれ以下の図の赤丸がある格子点, 青丸がある格子点へ行く通り数の和に対応します。
• 赤い格子点の一番下への行き方は \(\binom{N}{M}\) です。

以上より、上の式を二項係数に置き換えて

\[ f_M(N + 1) = 2 \times f_M(N) - \binom{N}{M}\]

という関係式を導出することができました。

よって二項係数を \( N \lt M \iff \binom{N}{M} = 0\) と定義することで、両方の場合を合わせて

\[f_M(N + 1) = 2 \times f_M(N) - \binom{N}{M}\]

と \(1\) 本の式で表すことができました。また、初項である

\[f_M(0) = 1\]

とあわせることで \(f_M(0), f_M(1), \dots, f_M(N)\) を \(\mathrm{O}(N)\) で列挙することもできました。

上の変形を用いれば、元の問題は \(i\) を昇順に見ていけば

\[ \sum_{a=0}^{\min(A,i)} \binom{i}{a} ,\sum_{b=0}^{\min(B,i)} \binom{i}{b}, \sum_{c=0}^{\min(C,i)} \binom{i}{c}\]

から

\[ \sum_{a=0}^{\min(A,i+1)} \binom{i+1}{a}, \sum_{b=0}^{\min(B,i+1)} \binom{i+1}{b}, \sum_{c=0}^{\min(C,i+1)} \binom{i+1}{c} \]

へ \(\mathrm{O}(1)\) で変形できることから、シグマ全体を \(\mathrm{O}(N)\) で計算できます。

\(1!,2!, \dots,N!\) と \(\frac{1}{1!}, \frac{1}{2!}, \dots, \frac{1}{N!}\) は \(\mathrm{O}(N + \log \mathrm{mod})\) あるいは \(\mathrm{O}(N)\) で列挙することができるので、この問題を\(\mathrm{O}(N + \log \mathrm{mod})\) あるいは \(\mathrm{O}(N)\) で解くことができました。

Python による実装例は以下の通りです。(AC するには処理系に PyPy3 を選ぶ必要があります。)

mod = 998244353
MAX = 5 * 10 ** 6 + 10

fac = [0] * (MAX + 1)
facinv = [0] * (MAX + 1)
fac[0] = 1
for i in range(1, MAX + 1):
  fac[i] = fac[i - 1] * i % mod
facinv[MAX] = pow(fac[MAX], mod - 2, mod)
for i in reversed(range(MAX)):
  facinv[i] = facinv[i + 1] * (i + 1) % mod

def binom(n, r):
  if n < 0 or r < 0 or n < r:
    return 0
  return fac[n] * facinv[r] % mod * facinv[n - r] % mod


N, A, B, C = map(int, input().split())
ans, a, b, c, sgn = 0, 1, 1, 1, 1 if N % 2 == 0 else -1
for i in range(N + 1):
  cur = binom(N, i) * a % mod * b % mod * c % mod
  ans += cur * sgn
  a = (a + a - binom(i, A)) % mod
  b = (b + b - binom(i, B)) % mod
  c = (c + c - binom(i, C)) % mod
  sgn = -sgn

print(ans % mod)

練習問題

この手の問題のように数え上げをグリッド上の頂点に言い換えるテクニックは実は ARC / AGC 典型ですが、 練習問題でこれらの問題を上げてしまうと重大なネタバレになってしまい大変良くないです。
そこで、グリッド上の斜め線が登場するテクニックを 非想定解法 として利用できる問題をいくつか挙げます。

1

(1) (解説の内容を踏まえれば易しいと思います)

\[f(N, M) = \begin{cases} 0 & \min(N,M) \lt 0 \\ \displaystyle\sum_{i=0}^{\min(N,M)} \binom{N}{i} & \mathrm{otherwise} \end{cases} \]

とおきます。 \(1!,2!, \dots,N!\) と \(\frac{1}{1!}, \frac{1}{2!}, \dots\) を事前に計算しているとき、 \(f(N,M)\) から \(f(N + 1, M), f(N - 1, M), f(N, M + 1), f(N, M - 1)\) を \(\mathrm{O}(1)\) で求めてください。

(2) ( (1) が分かれば 500~600 点相当です)

天下一プログラマーコンテスト2014本戦 D : 高橋君 を \(N = 100000\) として \(\mathrm{O}(N^{1.5})\) で解いてみましょう。

2

(600 点相当です) ABC234 F : Reordering の 文字種が \(2\) 種類の場合 を \(\mathrm{O}(|S|)\) で解いてみましょう。

3

(難しいです) Educational Codeforces Round 78 F : Cards を \(\mathrm{O}(N)\) で解いてみましょう。