誤読してタイムロスした。
https://yukicoder.me/problems/no/1075
問題
N頂点の木を成す無向グラフの各頂点に1~Kの値を振る。
この時、ある点vを1つ定めると、そこからどの点へのパスにおける頂点の値も、広義単調減少とできそうな振り方は何通りか。
解法
vを定めて、vから離れるほど単調減少になるケースを数え上げればいいのだが、複数同じ値の点がある場合、vの候補が複数あり得る場合がある。
その場合、最も親に近いものをvとするものとする。
dp(v,k) := vの頂点の値をkとし、vのsubtreeの頂点に対し、値が広義単調減少になるような値の振り方
dp2(v,k) := vの親頂点の値をk未満とし、v親方向の頂点群に対し、値が広義単調減少になるような 値の振り方
とすると、dp(v,k)*dp2(v,k)の総和が解。
dpは1回DFSを行って埋め、dp2はdpを用いて全方位DPの要領で親から求めて行こう。
int N,K; vector<int> E[1010]; const ll mo=1000000007; ll dp[1010][1010]; int C[1010]; ll ret; ll modpow(ll a, ll n = mo-2) { ll r=1;a%=mo; while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1; return r; } int dfs(int cur,int pre) { int i; C[cur]=1; FOR(i,K) dp[cur][i]=1; FORR(e,E[cur]) if(e!=pre) { C[cur]+=dfs(e,cur); ll sum=0; FOR(i,K) { (sum+=dp[e][i])%=mo; (dp[cur][i]*=sum)%=mo; } } return C[cur]; } void dfs2(int cur,int pre,vector<ll> P) { ll sum=0; int i; if(pre==-1) { FOR(i,K) ret+=dp[cur][i]; } else { FOR(i,K) { (ret+=dp[cur][i]*sum)%=mo; (sum+=P[i])%=mo; } sum=0; FOR(i,K) { (sum+=P[i])%=mo; (dp[cur][i]*=sum)%=mo; } } FORR(e,E[cur]) if(e!=pre) { sum=0; FOR(i,K) { (sum+=dp[e][i])%=mo; P[i]=dp[cur][i]*modpow(sum)%mo; } dfs2(e,cur,P); } } void solve() { int i,j,k,l,r,x,y; string s; cin>>N>>K; FOR(i,N-1) { cin>>x>>y; E[x-1].push_back(y-1); E[y-1].push_back(x-1); } dfs(0,0); vector<ll> P(K); dfs2(0,-1,P); cout<<ret%mo<<endl; }
まとめ
候補に迷うときは根に近い方に決める、というテクは応用ききそう。