kmjp's blog

競技プログラミング参加記です

AtCoder ABC #221 : H - Count Multiset

これ系さっと式が立たないな。
https://atcoder.jp/contests/abc221/tasks/abc221_h

問題

正整数N,Mが与えられる。
以下のmultisetは何通りか、各kについて求めよ。

  • 1~Nからなるk要素のmultisetである
  • 和はNである
  • 同じ値は最大M個しか含まれない

解法

f(n,s) := n要素からなり合計がsとなる、同じ値をM個以下しか含まないmultisetの組み合わせ
とすると、f(k,N)が解となる。

f(n,s)は以下のように求められる。

  • 1が含まれないケースを考えると、全要素から1を引いた状態から遷移できるのでf(n,s) += f(n,s-n)
  • 1が含まれるケースがある場合、1を取り除いた状態から遷移できる。ただし、すでにM個1を含んでいる場合は除くとすると、f(n,s) += f(n-1,s-1) - f(n-(M+1),s-n)

よってf(n,s)はO(ns)で求められる。

int N,M;
ll dp[5050][5050];
const ll mo=998244353;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	dp[0][0]=1;
	
	for(i=1;i<=N;i++) {
		for(j=i;j<=N;j++) {
			if(j>=i) dp[i][j]+=dp[i][j-i];
			dp[i][j]+=dp[i-1][j-1];
			if(i-1>=M&&j-1-(i-1)>=0) dp[i][j]+=mo-dp[i-1-M][j-1-(i-1)];
			dp[i][j]%=mo;
		}
		cout<<dp[i][N]<<endl;
	}
	
}

まとめ

色んな式変形の方法がありそう。