裏紙

ほぼ競プロ、たまに日記

CF 553A - Kyoya and Colored Balls

今回(やっと)Codeforcesに初参加した。1:30から2時間も考えてられず,70分くらいたったあたりで「もうやめて寝よう」ってなりました。最初の2問だけ解いて,この問題につまづいたので自分で文字に起こすことで理解を深めたい。こういう系の問題(というかDP)ってよく見るのになかなか自分はできないので,経験を積まねばと思う。これができるようになるともう一段階上がれると信じて頑張る。

・問題概要
Problem - 553A - Codeforces

異なるk個の色のボールが合計でn個ある(色をそれぞれ 色1,色2,...,色kと呼ぶとする)。
同じ色のボールは区別がつかない。
1個ずつ取り出していってn個のボールを並べる時,色iの最後の1個は色i+1の最後の1個よりも先に出る。(i=1,2,...k-1)

このようなボールの取り出し方は何通りあるか。1000000009で割った余りを出力。

k<=1000
n<=1000
色iのボールの個数c[i] (1<=c[i]<=1000)

・考え方
最後1個を別にとっておいて,それ以外を並べてからそのk色の最後の1個を入れる位置を決めればいいのかな?とか考えたけどそれだと(同じ色のボールを区別しないので)重複して数えちゃうやつとかがでてくる...?もうわかんないし寝よってなった。

やはりこういう問題はどこに視点を置くかが重要になる感じがある。何か1つに注目することが重要かもしれない。
f(i)を色1から色iまでの全てのボールを使って並べるときの解としよう。f(k)を求めたい。
まず,f(1)は(同色は区別しないので)1通り,f(1)=1である。
f(i)とf(i+1)の関係を考えよう。
既にi色の全てのボールが置かれている状況を考える。そこに,c[i+1]個の色i+1のボールを入れるとしよう。まず,条件から最後の1個は必ず色i+1である必要がある。残りの(c[i+1]-1)個はその列のどこに入れてもいい。つまり,c[1]+c[2]+...+c[i]+(c[i+1]-1)個のボールの並びの内,どの(c[i+1]-1)個の位置を色i+1の置き場にするかを選ぶということで,f(i+1)が決まる。

f(i+1) = f(i) * (c[1]+c[2]+...+c[i]+c[i+1]-1)C(c[i+1]-1)

組み合わせのCは,パスカルの三角形を使って事前に求めておく。これはO(n^2)。
問題を解くのにはO(k^2)。
これは間に合う。色を1個ずつ増やすという発想だった。

・実装(C++)

#include <cstdio>
#include <iostream>
using namespace std;

const long mod = 1000000007;

int main(){
  int k, ball[1000]={0};
  long c[1001][1001]={0};

  //input
  scanf(" %d", &k);
  for(int i=0; i<k; ++i) scanf(" %d", &ball[i]);

  //組み合わせの値をパスカルの三角形的に計算しておく
  for(int i=0; i<=1000; ++i){
    c[i][0]=1;
    c[i][i]=1;
  }

  for(int i=2; i<=1000; ++i){
    for(int j=1; j<i; ++j) c[i][j] = (c[i-1][j-1]+c[i-1][j])%mod;
  }

  long long ans=1;
  //solve
  for(int i=1; i<k; ++i){
    long first=-1;
    for(int j=0; j<=i; ++j) first+=ball[j];

    ans = (ans*c[first][ball[i]-1])%mod;
  }

  cout << ans << endl;
}