公式

G - At Most 2 Colors 解説 by en_translator


Let us first discuss the naivest solution.

One possible DP (Dynamic Programming) is \(dp[\) colors of last \(K\) squares \(] = \{\) number of ways \(\}\), but it costs like \(O(NC^{K+1})\), which is far from successful.

Let us discard unnecessary data from the key. The following information seems redundant:

  • when the last \(K\) squares contains three or more colors;
  • distinction of which colors are used.

Without them, we come up with the following DP: \(dp[\) the last position color \(X\) is used \(][\) the last position color \(Y\) is used \(]=\{\) number of cases \(\}\). However, it still costs like \(O(N^2)\), which don’t fit in the time limit.

There is one more parameter we can ignore: “the last position color \(Y\) is used.” One of the two colors is always used at the final square in the last \(K\) squares. Thus, when we are in square \(I\), one of the keys in the DP above is always \(i-1\), so we can eliminate one of the keys.

Now let us implement \(dp[\) the last position \(i\) of “the last color used, except for one used in the very last square (\(=\) the second last color appeared)” \(]=\{\) the number of ways \(\}\).
Note that if the last \(K\) squares are painted in the same color (that is, if you don’t need to consider the second last color appeared), you need to remember it using another variable \(sing\).

Initialize with \(sing=C\) and \(dp=(0,0,\dots,0)\), and do the transitions by determining the color of square \(i\):

  • Add \(dp[i-K]\) to \(sing\).
    • This is because \(dp[i-K]\) corresponds to the case where the \(K\) squares \(i-K,i-K+1,\dots,i-1\) is painted in the same color, and we do no longer need to maintain the information on the colors of the last \(K\) squares.
  • Let \(dp[i-1] \gets sing \times (C-1)\).
    • This corresponds to painting square \(i\) with a different color when the last \(K\) squares are painted in the same color.
    • This case, the rightmost occurrence of the second last color is at \(i-1\).
  • Add \(dp[i-1]\) to \(dp[i-K+1]+\dots+dp[i-2]\).
    • This corresponds to painting square \(i\) with the second last color; this updates the second last color.
    • This can be sped up using cumulative sums.
  • The transition itself is over.
    • You don’t need to reset \(sing\), and \(dp[i-2]\) and before. Keeping them as is corresponds to painting square \(i\) with the same color as square \(i-1\).

The time complexity is \(O(N)\).

Sample code (C++):

#include<bits/stdc++.h>
#include<atcoder/modint>

using namespace std;
using namespace atcoder;
using mint=modint998244353;

int main(){
  int n,k,c;
  cin >> n >> k >> c;
 
  vector<mint> dp(n+5,0);
  vector<mint> dp_sum(n+5,0);
  mint sing=c;
  
  for(int i=2;i<=n;i++){
    if(i-k>=0){ sing+=dp[i-k]; }

    dp[i-1]=(sing*(c-1));
    dp[i-1]+=dp_sum[i-2];
    if(i-k>=0){ dp[i-1]-=dp_sum[i-k]; }

    dp_sum[i-1]=dp_sum[i-2]+dp[i-1];
  }

  mint res=sing;
  for(int i=n-k+1;i<=n;i++){res+=dp[i];}
  cout << res.val() << "\n";
  return 0;
}

投稿日時:
最終更新: