kmjp's blog

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

天下一プログラマーコンテスト2019 : D - Three Colors

エクサウィザーズは良かったのにまたいまいち…。
https://atcoder.jp/contests/tenka1-2019/tasks/tenka1_2019_d

問題

N個の棒があり、各長さはA[i]とする。
各棒を3色のいずれかで塗り分けたのち、同じ色の棒を連結して3つの棒にする。
この時、この3つの棒の長さを持つような面積正の三角形が作れるような塗り分け方は何通りか。

解法

3色の棒の長さをそれぞれ状態として持とうとすると組み合わせが多くて破たんする。
逆に、前端の半分以上の長さを持つ辺があると三角形を作れないので、その組み合わせを求めよう。

f(n,l) := 最終的に最長となる棒の長さについて、n本目まで塗り分けた際の長さがlであるような組み合わせ
とする。(n+1)本目について、以下のように状態遷移できる。

  • 最長となる棒に追加するなら、f(n+1,l+A(n+1)) += f(n,l)
  • 残り2本のいずれかに追加するなら、f(n+1,A(n+1)) += 2*f(n,l)

最終的に、L≧N/2であるようなLに対し、f(N,L)*3の総和が解。(×3はどの色を最長にするかの組み合わせ)
ただ1点注意点があって、Nが偶数のとき、半分の長さの辺が2つできるケースが多重カウントされるので、これは別途計算する必要がある。

int N;
int A[303],S;
ll mo=998244353;

ll from[300*301][2];
ll to[300*301][2];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	from[0][0]=from[0][1]=1;
	ll ret=1;
	FOR(i,N) {
		cin>>A[i];
		S+=A[i];
		ret=ret*3%mo;
		ZERO(to);
		FOR(j,300*301) {
			if(from[j][0]) {
				// add
				(to[j+A[i]][0]+=from[j][0])%=mo;
				// non
				(to[j][0]+=from[j][0])%=mo;
			}
			if(from[j][1]) {
				// add
				(to[j+A[i]][1]+=from[j][1])%=mo;
				// non
				(to[j][1]+=2*from[j][1])%=mo;
			}
		}
		swap(from,to);
	}
	
	
	for(i=0;i<=S;i++) {
		if(i*2>=S) {
			ret-=from[i][1]*3;
		}
		if(i*2==S) {
			ret+=from[i][0]*3;
		}
		ret%=mo;
	}
	cout<<(ret+mo)%mo<<endl;
	
	
}

まとめ

これ思いつくのは速かったけど状態遷移を誤って地味に手間取った。