F - 1122 Subsequence 2 Editorial
by
sounansya
1122文字列 \(T\) の \(\displaystyle \frac{|T|}2\) 文字目を \(S\) の \(i\) 文字目に固定した場合の答えを求めることを考えます。 (この問題を部分問題と呼びます。)\(i=1,2,\ldots,|S|\) に対するこの値の総和がこの問題の答えとなることが分かります。
\(S\) の \(i\) 文字目より左にある数字 \(S_i\) の個数を \(p\) 、\(S\) の \(i\) 文字目より右にある数字 \(S_i+1\) の個数を \(q\) とします。この \(p,q\) は事前に累積和を計算することで高速に計算することができます。また、このとき部分問題の答えは \(\displaystyle \binom{p+q}{p+1}\) と表すことができます。
証明 \(1\):
1122文字列の長さを \(2k\) として固定すると、\(\displaystyle \frac{|T|}2\) 文字目が固定されていることから \(p\) 文字の中から \(k-1\) 文字、\(q\) 文字の中から \(k\) 文字選ぶ必要があります。したがって、部分問題の答えは
\[\begin{aligned} &\phantom{=} \sum_{k\geq 1} \binom{p}{k-1}\binom{q}k\\ &= \sum_{k\geq 1} \binom{p}{k-1}\binom{q}{q-k}\\ &=\sum_{k_1+k_2=q-1}\binom{p}{k_1}\binom{q}{k_2}\\ &=\left[x^{q-1}\right]\left(\sum_{k_1\geq 0}\binom{p}{k_1}x^{k_1}\right)\left(\sum_{k_2\geq 0}\binom{q}{k_2}x^{k_2}\right)\\ &=\left[x^{q-1}\right] (1+x)^{p+q}\\ &=\binom{p+q}{q-1}=\binom{p+q}{p+1} \end{aligned}\]
となります(この式変形はヴァンデルモンドの畳み込みと呼ばれます)。
証明 \(2\) :
\(p,q\) に数えられている添字の集合をそれぞれ \(X,Y\) とします(\(|X|=p,|Y|=q\) が成立します)。 \(X \cup Y\) からちょうど \(p+1\) 個の要素を選び、 \(X\) の要素のうち選ばれなかった添字、\(Y\) の要素のうち選ばれた添字、\(i\) を合わせた添字集合からなる文字列を \(T'\) とします。すると、 \(T'\) は 1122文字列 となり、さらに \(p+1\) 個の選び方と 1122文字列 は一対一に対応します。したがって、部分問題の答えは \(\displaystyle \binom{p+q}{p+1}\) となります。
以上より、 \(i=1,2,\ldots,N\) に対して上記の \(p,q\) を求め、求めた \(p,q\) に対し \(\displaystyle \binom{p+q}{p+1}\) の総和を計算すればこの問題を解くことができます。累積和や二項係数の前計算は \(O(|S|)\) 時間でできるので、全体の計算量も \(O(|S|)\) となります。
以上を適切に実装することでこの問題に正答することができます。
MOD = 998244353
INF = 2_000_000 + 2025
fact = [0] * (INF + 1)
finv = [0] * (INF + 1)
fact[0] = 1
for i in range(1, INF + 1):
fact[i] = fact[i - 1] * i % MOD
finv[INF] = pow(fact[INF], MOD - 2, MOD)
for i in range(INF, 0, -1):
finv[i - 1] = finv[i] * i % MOD
def comb(n, k):
if n < 0 or k < 0 or n < k:
return 0
return fact[n] * finv[k] % MOD * finv[n - k] % MOD
s = [ord(c) - ord('0') for c in input()]
n = len(s)
d = [[0] * (n + 1) for _ in range(10)]
for i in range(n):
for j in range(10):
d[j][i + 1] = d[j][i]
d[s[i]][i + 1] += 1
ans = 0
for i in range(n - 1):
if s[i] == 9:
continue
p = d[s[i]][i + 1] - 1
q = d[s[i] + 1][n] - d[s[i] + 1][i + 1]
ans = (ans + comb(p + q, p + 1)) % MOD
print(ans)
posted:
last update:
