안 어려운 문제였는데 루트를 바꾼다는 생각을 못해서 1시간이나 써버렸어....

난 왜 트리의 루트를 1번으로 고정해놓고 생각할까


  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. typedef long long ll;
  4. vector<int>tree[200005];
  5. bool task[200005];
  6. ll DP[200005],dist[200005];
  7. int depth[200005];
  8. int dp[200005][20];
  9. void DFS(int cur,int p){
  10. DP[cur] = task[cur];
  11. for(int nxt:tree[cur]){
  12. if(nxt==p) continue;
  13. depth[nxt] = depth[cur]+1;
  14. dp[nxt][0] = cur;
  15. DFS(nxt,cur);
  16. DP[cur]+=DP[nxt];
  17. if(DP[nxt]) dist[cur]+=dist[nxt]+2;
  18. }
  19. }
  20. int LCA(int a,int b){
  21. if(depth[a]<depth[b]) swap(a,b);
  22. int d = depth[a] - depth[b];
  23. for(int i=0; d; i++){
  24. if(d&1) a = dp[a][i];
  25. d>>=1;
  26. }
  27. if(a!=b){
  28. for(int i=19; i>=0; i--){
  29. if(dp[a][i]==dp[b][i]) continue;
  30. a = dp[a][i]; b = dp[b][i];
  31. }
  32. a = dp[a][0];
  33. }
  34. return a;
  35. }
  36.  
  37. void solve(){
  38. int n,k; cin>>n>>k;
  39. int x,y; cin>>x>>y;
  40. for(int i=1; i<=n; i++){
  41. task[i] = depth[i] = DP[i] = 0;
  42. dist[i] = 0;
  43. tree[i].clear();
  44. }
  45. task[x] = task[y] = 1;
  46. while(k--){
  47. int c; cin>>c;
  48. task[c] = 1;
  49. }
  50. k = 0;
  51. for(int i=1; i<=n; i++) k+=task[i];
  52. for(int i=1; i<n; i++){
  53. int a,b; cin>>a>>b;
  54. tree[a].emplace_back(b);
  55. tree[b].emplace_back(a);
  56. }
  57. DFS(1,0);
  58. for(int i=1; i<20; i++){
  59. for(int j=1; j<=n; j++) dp[j][i] = dp[dp[j][i-1]][i-1];
  60. }
  61. int lca = LCA(x,y);
  62. ll ans = 0;
  63. int pre = 0;
  64. if(lca==x || lca==y){
  65. if(lca==y) swap(x,y);
  66. while(y!=lca){
  67. ans+=dist[y];
  68. if(pre) ans-=dist[pre]+2;
  69. pre = y;
  70. ans++;
  71. y = dp[y][0];
  72. }
  73. ans+=dist[x];
  74. if(pre) ans-=dist[pre]+2;
  76. while(x!=1){
  77. if(DP[x]<k){
  78. ans+=2;
  79. int p = dp[x][0];
  80. ans+=dist[p] - (dist[x]+2);
  81. x = p;
  82. }
  83. else break;
  84. }
  85. }
  86. else{
  87. int pre = 0;
  88. while(x!=lca){
  89. ans+=dist[x];
  90. if(pre) ans-=dist[pre]+2;
  91. pre = x;
  92. ans++;
  93. x = dp[x][0];
  94. }
  95. int PRE = 0;
  96. while(y!=lca){
  97. ans+=dist[y];
  98. if(PRE) ans-=dist[PRE]+2;
  99. PRE = y;
  100. ans++;
  101. y = dp[y][0];
  102. }
  103. ans+=dist[lca];
  104. if(pre) ans-=dist[pre]+2;
  105. if(PRE) ans-=dist[PRE]+2;
  106. while(x!=1){
  107. if(DP[x]<k){
  108. ans+=2;
  109. int p = dp[x][0];
  110. ans+=dist[p] - (dist[x]+2);
  111. x = p;
  112. }
  113. else break;
  114. }
  115. }
  116. cout<<ans<<'\n';
  117. }
  118. int main(){
  119. ios_base::sync_with_stdio(false); cin.tie(NULL);
  120. int t = 1;
  121. cin >> t;
  122. while(t--) solve();
  123. }