SRM564 Div1 Medium(500) AlternateColors2
素朴に求めるプログラムを書いて、その結果から一般項を推測した。
n=kの場合はn個のボールを赤が最多となるように塗り分ける場合の数。kを固定した場合の答えは、階差数列が等差数列になっている。
long long naive( int n, int k ) { long long ans = 0; for ( int R=0; R<=n; R++ ) for ( int G=0; R+G<=n; G++ ) { int r = R; int g = G; int b = n-R-G; int c = 0; while ( r+g+b>0 ) { if ( r>0 ) { r--, c++; if ( c==k ) ans++; } if ( g>0 ) g--, c++; if ( b>0 ) b--, c++; } } return ans; } class AlternateColors2{public: long long countWays( int n, int k ) { if ( k==1 ) return (long long)n*(n+1)/2; long long ans = 0; // n==kの場合を求める // k個のボールを赤が最多になるように塗り分ける場合の数 for ( int r=(k+4)/3; r<=k; r++ ) if ( k-r <= r-1 ) ans += k-r+1; else ans += 2*(r-1)-(k-r)+1; long long s; long long d; if ( (k+2)%3==0 ) { s = ((k+2)/3+1)/2*2; d = 1; } else { s = (k+3)/6*2; d = 0; } for ( int i=k; i<n; i++ ) { ans += s; s += d; } return ans; }};