[ABC222] E - Red and Blue Tree

[Q] https://atcoder.jp/contests/abc222/tasks/abc222_e

・クエリMが100、頂点数Nが1000と小さいので、この処理は毎回やっても、O(100 * 1000)でネックにならなそう。
・辺ごとに通過した回数をおさえておけば、通過回数の組み合わせの探索で答えが求まりそう。

Q, 通過回数？

A. 入力例1
4 5 0
2 3 2 1 4
1 2
2 3
3 4

・辺 1-2 は 2回 とおる。
Edge[0] = 2
・辺 2-3 は 3回 とおる。
Edge[1] = 3
・辺 3-4 は 1回 とおる。
Edge[2] = 1

総通過数 = Edge[0 ~ 2] = 2+3+1 = 6
K = 0なので、
R=3, B=3 の取り方を満たす組み合わせが回答になる。

幸い、Edgeの総数の上限が、1000 * 100 = 100000までとわかるので、
DP[ M-1 ][ 100000 ]  = 組み合わせ数が間に合うことに気づく。
解答は、R値の組み合わせ数が分かればいい。

Edgeの総数=10
K=2
R-B=2を満たすのは6-4
(10 - K) / 2 の差分が、2で割り切れることが解答を持つ条件になりそう。
Ｒは、K + (10 - K) / 2 で求まる。

解答はDP[ M-1 ][ sum-K)/2 + K ]。

実装。

ll DP[100010];
ll pid[1010][1010]; // 辺id。pid[A][B] A->Bの辺
ll cnts[1010]; // cnts[辺id] その辺を何回通過したか

int main(){
 cincout();
 
 ll N, M, K;
 cin >> N >> M >> K;
 graph g(N+1); // 0index ~ 
 // 戻り道
 ll A[M];
 rep(i, M) cin >> A[i];
 rep(i, M) --A[i];
 rep(i, N-1){
   ll U, V;
   cin >> U >> V;
   --U, --V;
   g.add_edge(U, V, 1);
   g.add_edge(V, U, 1);
   pid[U][V] = i;
   pid[V][U] = i;
 }
 bool checked[N]={};
 rep(i, M-1){
   ll start = A[i];
   ll goal = A[i+1];
   g.dijkstra(start);
   vector<ll> path = g.get_path(goal, start);
   ll psz = path.size();
   rep(j, psz-1){
     ll ps=path[j];
     ll pg=path[j+1];
     ++cnts[ pid[ps][pg] ]; // 
   }
 }
 DP[0]=1;
 ll sum = 0;
 rep(i, N-1){
   sum += cnts[i];
   for(ll j=100000; j>=cnts[i]; --j){
     ( DP[j] += DP[j - cnts[i]] ) %= MOD;
   }
 }
 ll ans = sum-K;
 if (ans%2){
   cout << 0 << endl;
   return 0;
 }
 cout << DP[ans/2+K] << endl;
}

Q. dijkstra？
A. ダイクストラというより、get_pathなる便利なものを仕込んでいたので使ったまで。
dfsで、どのEdgeを通過したかさえ得られればもっと小さい実装で軽くいけそう。

https://atcoder.jp/contests/abc222/submissions/26456103