公式

F - Second Gap 解説 by harurun4635


naive な DP

順列の後ろから順に要素を追加することを考えましょう。

\((P_{i+1}, \dots, P_{N})\) の先頭に新しく \(P_i\) を追加する事を考えます。ただし、このとき「相対的な大小」のみを考えます。(いわゆる挿入 DP とおなじ考え方です)

いまの max, second-max の index を \(\{a, b\}\) のように表すことにします。

このとき、 \((P_i, P_{i+1}, \dots, P_N)\) における \(P_i\) の大きさに応じて、以下のように遷移します。

  • \(P_i\)\(1\) 番目に大きいときは \(\{i, a\}\)

  • \(P_i\)\(2\) 番目に大きいときは \(\{a, i\}\)

  • それ以外のときは \(\{a,b\}\)

これを DP のキーに持ち、遷移ごとに \(D_i\) の条件を満たすか確認することで、多項式時間の解法ができます。

key の削減

実は \(b\) を管理しなくても DP ができます。

まず \(P_i\)\(1\) 番目か \(2\) 番目になるとき、 \(b\) は関係ないです

また、それ以外の遷移であっても \(\{a, b\}\) のままですから \(b\) に着目せずとも、 \(D_i = D_{i+1}\) である時に限り条件を自動的に満たすことがわかります。

ということで、実は \(b\) を管理しなくてもよいです。その上で、条件を整理して DP を書き直しましょう。

DP

\(dp[i][a] =\) 最大値が \(P_a\) であるような \((P_i, \dots, P_N)\) の(相対的な大小の)通り数

とします。初期化は \(dp[n-1][n-1] = dp[n-1][n] = 1\) とします。

\(i\) の降順に遷移は以下のとおりにすればよいです。

  • \(P_i\)\(1\) 番目になるとき

    直前の最大値の位置 \(a\)\(|a - i| = D_i, a > i\) ですから \(a = i + D_i\) です。
    よって \(dp[i][i] += dp[i+1][i + D_i]\) と遷移できます。

  • \(P_i\)\(2\) 番目になるとき

    直前の最大値の位置は同様に \(a = i + D_i\) です。
    よって \(dp[i][i+D_i] += dp[i+1][i+D_i]\) と遷移できます。

  • それ以外のとき

    まず、それ以外の個数は \(N-1-i\) 通りあります。
    そして、先ほど述べた通り \(D_i = D_{i+1}\) なのであれば、すべての \(j\) に対して
    \(dp[i][j] += dp[i+1][j]\) と遷移すればよいです。

高速化

このままでは、\(i\) ごとに遷移が \(O(N)\) です。

しかし、一番下の遷移は全体に定数をかけるのみです。これを除けば遷移は \(i\) ごとに \(O(1)\) となります。
よって DP を in-place に行い、「今かけられている定数」を DP 配列の外側で管理することで \(O(N)\) で実装できます。(くわしくは実装例を参照してください)

また、双対セグ木や遅延セグ木などを用いて \(O(N \log N)\) でも簡単でしょう。

実装例

簡単のため辞書型を用い、逆元の計算を \(O(\log MOD)\) で行っています。

from collections import defaultdict

mod = 998244353

n = int(input())
d = list(map(int, input().split()))

dp = defaultdict(int)
dp[n-1] = 1
dp[n-2] = 1
coef = 1  # かけられた定数を管理する

for i in reversed(range(0, n - 2)):
    j = i + d[i] # 直前の最大値の位置
    v = dp[j]
    if d[i] == d[i + 1]:
        coef = coef * (n - i - 2) % mod # 3 番目の遷移
        add = v * pow(n - i - 2, -1, mod) % mod # かけられた定数の分補正をする
        dp[j] = (dp[j] + add) % mod # 2 番目の遷移
        dp[i] = (dp[i] + add) % mod # 1 番目の遷移
    else:
        dp.clear() # 3 番目の遷移 in-place なので削除をする
        dp[j] = v # 2 番目の遷移
        dp[i] = v # 1 番目の遷移

ans = sum(dp.values()) % mod
print(ans * coef % mod)

以下は遅延セグ木を用いた解法です。 加算が \(1\) 点しかないことを利用していることに注意してください。

from atcoder.lazysegtree import LazySegTree

mod = 998244353

n = int(input())
d = list(map(int, input().split()))

def add(x, y): return (x + y) % mod
def mul(x, y): return x * y % mod

dp = LazySegTree(add, 0, mul, mul, 1, [0] * (n - 2) + [1] * 2)

for i in reversed(range(0, n - 2)):
    j = i + d[i]
    v = dp.get(j)
    if d[i] == d[i+1]:
        dp.apply(0, n, n - i - 2)
    else:
        dp.apply(0, n, 0)
    dp.set(j, (dp.get(j) + v) % mod)
    dp.set(i, (dp.get(i) + v) % mod)

print(dp.all_prod())

投稿日時:
最終更新: