dp による解法を紹介します。

累積和によって、次の問題に言い換えておきます。

\(0\), \(1\), \(2\) からなる数列 \((A_0, A_1, \ldots, A_N)\) がある。\(A_0\) と \(A_N\) を含む \(K+1\) 項からなる部分列であって、隣接する \(2\) 項がすべて異なるものを見つけよ。

まず、次を計算する dp が考えられそうです。

• \(dp[k][i]\) (\(1\leq k\leq N+1, 0\leq i\leq N\))：\(A_i\) を最後に選ぶ部分列であって、項数が \(k\) のものが存在するかの bool 値

これでは状態数が \(\Theta(N^2)\) となってしまって解くのは難しいです。

この dp には、次の単調性があります。

• \(0 < i < j, A_i = A_j\) かつ \(dp[k][i]\) が true ならば、 \(dp[k][j]\) も true である。

この単調性は、最後に選んだ項 \(A_i\) を \(A_j\) に置き換えることで分かります。これによって、計算対象を次の \(O(N)\) 状態に限定できます。

• \(dp[k][x]\) (\(1\leq k\leq N+1, 0\leq x\leq 2\))：\(A_i = x\) を最後に選ぶ部分列で、項数が \(k\) になるものが存在するような、最小の \(i\)

\(dp[k][-]\) は \(dp[k-1][-]\) から簡単に計算できます。あとは \(dp\) テーブルから解を復元すればよいです。

• 解答例：https://atcoder.jp/contests/jsc2022-final/submissions/34330925

この実装は \(\Theta(N\log N)\) 時間での解法ですが、簡単に \(O(N)\) 時間にできます。