SRM564 Div1 Medium(500) AlternateColors2

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;
}};