akashtadwai

Painting-Fence (DP)

Sep 22nd, 2021
700
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. long long countWays(int n, int k){
  2. // dp[i][j][k] -> number of ways from first i elements
  3. // where ith color is j and there are k consecutive elements of color j
  4. // dp[i][j][2] = dp[i-1][j][1]
  5. // dp[i][j][1] = sigma_{x!=j} dp[i-1][x][{1,2}]
  6. int dp[n+1][k+1][3];
  7. memset(dp,0,sizeof(dp));
  8. for(int i=1;i<=k;i++) dp[1][i][1] = 1;
  9. for(int i=2;i<=n;i++){
  10.     ll s= 0 ;
  11.     for(int j=1;j<=k;j++){
  12.         add_self(s,dp[i-1][j][1]);
  13.         add_self(s, dp[i-1][j][2]);
  14.     }
  15.     for(int j=1;j<=k;j++){
  16.         ll x = s;
  17.         dp[i][j][2] =  dp[i-1][j][1];
  18.         sub_self(x, dp[i-1][j][1]);
  19.         sub_self(x,dp[i-1][j][2]);
  20.         dp[i][j][1]=x;
  21.     }
  22. }
  23. ll ans=0;
  24. for(int i=0;i<=k;i++) {
  25.     for(int j=0;j<=2;j++)
  26.         add_self(ans,dp[n][i][j]);
  27. }
  28. return ans;
  29.     }
RAW Paste Data