わりかし面倒な問題。
https://yukicoder.me/problems/no/1502
問題
長さNの正整数列Aに対し、以下を満たすものは何通りか。
- Aの最大値はK
- A[i]+A[j]=***といった2要素の和に関する条件が与えられる
解法
N要素のグラフを考える。条件は、すなわち重み付きの辺とする。
この際、連結成分毎に考える。
- 連結成分が二部グラフである場合
- 1要素を決めると他の要素も決まる。その1要素の上限・下限を定め、以後DFSしておくことで、最大値がK以下という条件から最初の要素の上限・下限を更新し求めることができる。
- 連結成分が二部グラフでない場合
- 値が一意に定まるので、最大値の条件を確認する
上記結果を連結成分毎に掛け合わせていくわけだが、最大値がK未満の場合とKの場合をそれぞれ数え上げることで、最大値がKの連結成分が1個以上含まれるようにする。
int N,M,K; vector<pair<int,int>> E[101010]; const ll mo=1000000007; ll num[2][101010]; set<int> cand; void dfs(int cur,int v,int id) { cand.insert(cur); if(num[id][cur]!=-1LL<<60) { if(num[id][cur]!=v) { cout<<0<<endl; exit(0); } return; } num[id][cur]=v; FORR2(e,c,E[cur]) dfs(e,c-v,id^1); } void solve() { int i,j,k,l,r,x,y; string s; cin>>N>>M>>K; K--; FOR(i,M) { cin>>x>>y>>k; k-=2; if(k<0) return _P("0\n"); E[x-1].push_back({y-1,k}); E[y-1].push_back({x-1,k}); } FOR(i,N) num[0][i]=num[1][i]=-1LL<<60; ll from[2]={1,0}; FOR(i,N) if(num[0][i]==-1LL<<60&&num[1][i]==-1LL<<60) { ll to[2]={0,0}; cand.clear(); dfs(i,0,0); if(num[1][i]==-1LL<<60) { ll L=0,R=K; FORR(c,cand) { ll a=num[0][c]; ll b=num[1][c]; if(a!=-1LL<<60) { L=max(L,-a); R=min(R,K-a); } else { L=max(L,b-K); R=min(R,b); } } if(L>R) { cout<<0<<endl; return; } ll v=(R-L+1); set<int> S; FORR(c,cand) { if(num[0][c]+R==K) S.insert(R); if(num[1][c]-L==K) S.insert(L); } to[0]=from[0]*(v-S.size())%mo; to[1]=(from[0]*S.size()+from[1]*v)%mo; } else { int ok=1; int ma=0; FORR(c,cand) { ll v=num[0][c]+num[1][c]; if(v%2||v<0) { cout<<0<<endl; return; } v/=2; if(v==K) ma=1; } to[ma]=(from[0]+from[1])%mo; } swap(from,to); } cout<<from[1]<<endl; }
まとめ
方針は立っても、実装詰めていくのが大変。