Naive DP (Dynamic Programming) solution

Let us construct a permutation by repeatedly prepending an element.

Suppose we are now inserting \(P_i\) to \((P_{i+1}, \dots, P_{N})\). Here, we only care the relative magnitude of the values. (The idea is similar to Insertion DP.)

Let \(\{a, b\}\) denote the indices of the maximum and second maximum values before the insertion.

Depending on the magnitude of \(P_i\) among \((P_i, P_{i+1}, \dots, P_N)\), it becomes:

• \(\{i, a\}\) if \(P_i\) is the largest;

• \(\{a, i\}\) if \(P_i\) is the second largest;

• \(\{a,b\}\) otherwise.

By maintaining them as a DP key and checking if the condition of \(D_i\) is satisfied on each transition step, we obtain a polynomial-time solution.

Reducing keys

In fact, we may perform DP without maintaining \(b\).

First, if \(P_i\) becomes largest or second largest, \(b\) is irrelevant.

For the other case, it remains \(\{a, b\}\), so the condition is automatically satisfied only when \(D_i = D_{i+1}\); again, \(b\) does not matter.

Therefore, we do not have to memorize \(b\) at all. Let us reorganize the condition and rewrite the DP.

The DP

Let \(dp[i][a] =\) the number of (relative magnitudes of) \((P_i, \dots, P_N)\) whose maximum value is \(P_a\).

The initial values are \(dp[n-1][n-1] = dp[n-1][n] = 1\).

The table can be filled in descending order of \(i\), according to the following transitions:

• If \(P_i\) becomes largest

  The position \(a\) of the largest element before the insertion is \(|a - i| = D_i, a > i\), so \(a = i + D_i\).
  Thus, the transition is \(dp[i][i] += dp[i+1][i + D_i]\).

• If \(P_i\) becomes second largest

  The position of the largest element before the insertion is \(a = i + D_i\).
  Thus, the transition is \(dp[i][i+D_i] += dp[i+1][i+D_i]\).

• Otherwise

  There are \((N-1-i)\) cases that fall into this pattern.
  As described before, if \(D_i = D_{i+1}\), we may transition as
  \(dp[i][j] += dp[i+1][j]\) for all \(j\).

Optimization

Here, the transition costs \(O(N)\) time for each \(i\).

However, the final transition is merely multiplying a constant, and all other transitions is \(O(1)\).
Therefore, we can perform the DP in-place, and manage the “current constant factor” outside the DP table. This way, the DP runs in a total of \(O(N)\) time. (Please refer to the sample code for details.)

Alternatively, one may easily implement it using a dual segment tree or a lazy segment tree.

Sample code

For simplicity, we use the dictionary type, and spend \(O(\log MOD)\) time to compute an inverse.

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  # Manages the global factor

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 # Transition for case 3
        add = v * pow(n - i - 2, -1, mod) % mod # Correct for the global factor
        dp[j] = (dp[j] + add) % mod # Transition for case 2
        dp[i] = (dp[i] + add) % mod # Transition for case 1
    else:
        dp.clear() # Transition for case 3; clear elements because this is in-place DP
        dp[j] = v # Transition for case 2
        dp[i] = v # Transition for case 1

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

This is a solution using a lazy segment tree. Note that we only need to apply element-wise addition.

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())