😎

ABC 458 E - Count 123

に公開
競技プログラミング
AtCoder
競プロ
アルゴリズム
tech

ABC 458 E - Count 123

問題

問題概要

以下のような数列の個数を数え上げる。

  • i を X_i 個含む。(1 \leq i \leq 3)
  • 隣接する要素の差は高々 1

今回の問題のポイント

  1. 2 を固定して考える。
  2. 2 の前後にできる各隙間には、以下のいずれかを配置できる。
    • 1 を 1 個以上
    • 3 を 1 個以上
    • 何も置かない
  3. 組み合わせを効率的に数え上げる。

うまく行った方法

1 番目の条件だけなら組み合わせの問題なので簡単。
難しいと感じたのは 2 番目の条件だったので 2 番目の条件について深掘り。

2 番目の条件について深掘り

1, 2, 3 のそれぞれの要素について、次に来る要素を並べてみると、2 については 1,2,3 のどれでもいいので実質制約がないことがわかる。
1 と 3 は自分自身か 2 だけが隣接できる。

2 を先に並べると、2 の前後に合計 X_2+1 個の隙間ができる。

例えば、X_2=3 のときは以下のようになる。

_ 2 _ 2 _ 2 _

ここで重要なのは、同じ隙間に 1 と 3 を両方置くことはできないという点である。
なぜなら、同じ隙間に 1 と 3 が入ると、どこかで 1 と 3 が隣接してしまい、隣接差が 2 になって条件を満たさないからである。

したがって、各隙間は次のいずれかになる。

  • 1 だけを置く隙間
  • 3 だけを置く隙間
  • 何も置かない隙間

そこで、まず 1 を置く隙間の数を n_1 として固定する。
1 を置かない残りの隙間の数を n_3=X_2+1-n_1 とする。

1 を置く隙間は空であってはいけない。
一方で、残りの n_3 個の隙間には 3 を 0 個以上置けばよい。

解答方針

  • C_1: 1をおく隙間の選び方
C_1:=\binom{X_2+1}{n_1}
  • C_2: 1をn_1個の隙間に置く置き方(各1つ以上)
C_2:=\binom{X_1-1}{n_1-1}
  • C_3: 3をn_3個の隙間に置く置き方(各0個以上)
C_3:=\binom{X_3+n_3-1}{n_3-1}

となる。
特に C_2, C_3 は写像12相から導かれる。それぞれ

  • C_2 は X_1 個の区別しない玉を n_1 個の区別する箱に各1つ以上入れるパターン
  • C_3 は X_3 個の区別しない玉を n_3 個の区別する箱に各0個以上で入れるパターン

に該当する。

あとは、競プロの典型だが、毎回二項係数を一から計算すると時間が足りない。
今回は逆元を前計算し、n_1 を 1 つ増やすたびに C_1, C_2, C_3 を更新することで効率的に求める。

解答例

MOD = 998244353

X1, X2, X3 = map(int, input().split())

max_n = 2 * max(X1, X2, X3) + 1
inv = [0] * (max_n + 1)
inv[1] = 1
for i in range(2, max_n + 1):
    inv[i] = MOD - (MOD // i) * inv[MOD % i] % MOD


ans = 0
comb_choice_1 = 1  # C1: 1 を入れる隙間の選び方
comb_put_1 = 1  # C2: 1 を選んだ隙間に入れる場合の数
comb_put_3 = 1  # C3: 3 を残りの隙間に入れる
# n1 = 0 のとき、3 を入れてよい隙間は X2 + 1 個。
# comb_put_3 = C(X3 + (X2 + 1) - 1, (X2 + 1) - 1)
#            = C(X3 + X2, X2)
# を先に作っておく。
for i in range(1, X2 + 1):
    comb_put_3 *= X3 + i
    comb_put_3 *= inv[i]
    comb_put_3 %= MOD

for n1 in range(1, min(X1, X2) + 1):
    # n1: 1 を入れる 2 の隙間の数
    # n3: 3 を入れて良い 2 の隙間の数
    n3 = X2 + 1 - n1

    comb_choice_1 *= X2 + 1 - (n1 - 1)
    comb_choice_1 *= inv[n1]
    comb_choice_1 %= MOD

    comb_put_3 *= n3
    comb_put_3 *= inv[X3 + n3]
    comb_put_3 %= MOD

    ans += comb_choice_1 * comb_put_1 * comb_put_3
    ans %= MOD

    comb_put_1 *= X1 - n1
    comb_put_1 *= inv[n1]
    comb_put_1 %= MOD

print(ans)

次回以降に生かすポイント

  • 不変な部分を探して軸にして考える
    • 今回は、自由に置くのが難しい 1 と 3 ではなく、どの要素とも隣接できる 2 を先に固定した。
  • 「隙間に分配する」形に持ち込めないか考える
    • 並べ方の問題は、基準となる要素を固定すると、区別する箱への分配として数えられることがある。
  • 重複して数えていないか確認する
    • 今回は、1 を置く隙間を「空なし」で固定することで、同じ数列を複数回数えないようにした。
GitHubで編集を提案

Discussion