kmjp's blog

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

yukicoder : No.1227 I hate ThREE

なぜこれに気が付かなかったのか。
https://yukicoder.me/problems/no/1227

問題

木を成すN頂点の無向グラフが与えられる。
各頂点に1~Cの整数を振ったとき、隣接点に振った整数の差が3となるのは何通りか。

解法

適当な頂点を根として考える。
もしCが6Nより大きいとき、根頂点の値が3N~(C-3*N+1)の場合、そこからの隣接点の値が+3しようが-3しようが1~Cの範囲からはみ出ることはない。
よって、このようなケースは(C-6N)*2^(N-1)だけ先に解に計上しておこう。
あとはCが6Nであると考えればよい。

CがO(N)なら、O(N^2)のDPで愚直に計算して間に合う。

int N,C;
int S[3];
vector<int> E[1010];
int V[1010];
ll dp[1010][6010];
const ll mo=1000000007;

void dfs(int cur,int pre) {
	int i;
	FOR(i,C) dp[cur][i]=1;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		
		FOR(i,C) {
			ll cand=0;
			if(i-3>=0) cand+=dp[e][i-3];
			if(i+3<C) cand+=dp[e][i+3];
			dp[cur][i]=dp[cur][i]*cand%mo;
		}
	}
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>C;
	ll ret=0;
	if(C>6*N) {
		ret=C-(6*N);
		C=6*N;
		FOR(i,N-1) ret=ret*2%mo;
	}
	
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,0);
	FOR(i,C) ret+=dp[0][i];
	cout<<ret%mo<<endl;
	
}

まとめ

組み合わせや包除原理にとらわれすぎた。