部分列 DP メモ - noshi91のメモ の解法を改造して解きます。

[1] DP 定義

元の記事と同様に、次のように DP を定義します。添字は \(0\)-indexed とします。

• \(\mathrm{dp}[i][c] \coloneqq A_0, \dots, A_{i-1}\) の部分列であって 221 数列であるもののうち、\(c\) で終わるものの種類数
• \(\mathrm{sum}[i] \coloneqq A_0, \dots, A_{i-1}\) の部分列であって 221 数列であるものの種類数（空列含む）

すべての \(i\) について \(\displaystyle \mathrm{sum}[i] = 1 + \sum_c \mathrm{dp}[i][c]\) が成り立つことに注意します。

[2] DP 遷移

各 \(i\) に対して、\(i\) から左に見ていったときに（自身を含めて）\(A_i\) 個目に値 \(A_i\) が現れる位置を \(\mathrm{prv}[i]\) とします（存在しない場合の扱いは省略します）。遷移は次のようになります。

• \(A_i \neq c\) に対し、\(\mathrm{dp}[i+1][c] = \mathrm{dp}[i][c]\) です。
• \(\mathrm{dp}[i+1][A_i]\) について、\(A_i\) で終わる部分列は直近 \(A_i\) 個の \(A_i\) を使用するものを数えることにしてよいです。すると数えるべきものは \(A_0, \dots, A_{\mathrm{prv}[i]-1}\) の部分列であって 221 数列であるもののうち \(A_i\) 以外で終わるものの種類数であるので、\(\displaystyle \mathrm{dp}[i+1][A_i] = 1 + \sum_{c \neq A_i} \mathrm{dp}[\mathrm{prv}[i]][c] = \mathrm{sum}[\mathrm{prv}[i]] - \mathrm{dp}[\mathrm{prv}[i]][A_i]\) となります。
• \(\displaystyle \mathrm{sum}[i+1] = 1 + \sum_c \mathrm{dp}[i+1][c]\) を式変形すると、\(\mathrm{sum}[i+1] = \mathrm{sum}[i] + \mathrm{sum}[\mathrm{prv}[i]] - \mathrm{dp}[i][A_i] - \mathrm{dp}[\mathrm{prv}[i]][A_i]\) となります。

[3] 高速化

この DP を高速に計算することを考えます。

\(\mathrm{dp}[i][c]\) については、\(i\) を陽に持たずに in-place に更新していけばよいです。

\(\mathrm{dp}[\mathrm{prv}[i]][A_i]\) を取得する必要があることについては、\(A_i = A_{\mathrm{prv}[i]}\) であることに注目すると、\(\mathrm{mem}[i] \coloneqq \mathrm{dp}[i][A_i]\) を適宜メモしておけばよいことがわかります。なおこれに気づかなくても、部分永続配列 を用いて処理してしまうことも可能です。

すべてを適切に実装すると、時間計算量は \(O(N)\) です。

C++ による実装例：

#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;

int main()
{
  int N;
  cin >> N;
  vector<int> A(N);
  for (int i = 0; i < N; i++)
    cin >> A[i];

  vector<vector<int>> I(N + 1);
  for (int i = 0; i < N; i++)
    I[A[i]].push_back(i);
  vector<int> prv(N, -1);
  for (int a = 1; a <= N; a++)
  {
    for (int k = a - 1; k < (int)I[a].size(); k++)
    {
      prv[I[a][k]] = I[a][k - (a - 1)];
    }
  }

  vector<mint> dp(N + 1, 0), sum(N + 1, 0), mem(N + 1, 0);
  sum[0] = 1;
  for (int i = 0; i < N; i++)
  {
    int a = A[i], j = prv[i];
    mem[i] = dp[a];
    dp[a] = (j >= 0 ? sum[j] - mem[j] : 0);
    sum[i + 1] = sum[i] - mem[i] + dp[a];
  }
  cout << (sum[N] - 1).val() << endl;
}