fork download
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3.  
  4. #define ll long long
  5. #define sp ios_base::sync_with_stdio(false),cin.tie(NULL),cout.tie(NULL)
  6. #define cps CLOCKS_PER_SEC
  7. #define mod (ll)1000000007
  8. #define f first
  9. #define s second
  10. #define debug1(x) cout<<x<<"\n"
  11. #define debug2(x,y) cout<<x<<" "<<y<<"\n"
  12. #define debug3(x,y,z) cout<<x<<" "<<y<<" "<<z<<"\n"
  13. #define nl cout<<"\n";
  14. #define pq priority_queue
  15. #define inf 0x3f3f3f3f
  16. #define test cout<<"abcd\n";
  17. #define pi pair<int,int>
  18. #define pii pair<int,pi>
  19. #define pb push_back
  20. #define mxn 100005
  21. ll ans=0,d;
  22. int a[mxn],sz[mxn];
  23. map<ll,ll> mp[mxn];
  24.  
  25. void init(){
  26. for(int i=1; i<mxn; ++i){
  27. a[i]=i;
  28. sz[i]=1;
  29. }}
  30.  
  31. int root(int i){
  32. while(i!=a[i]){
  33. a[i]=a[a[i]];
  34. i=a[i];
  35. }
  36. return i;
  37. }
  38.  
  39. void unn(int p,int q){
  40. p=root(p);
  41. q=root(q);
  42. if(p==q)return;
  43. if(sz[p]<sz[q])swap(p,q);
  44. a[q]=a[p];
  45. sz[q]=0;
  46. sz[p]+=sz[q];
  47. for(auto &it:mp[q])
  48. for(int i=-d; i<=d; ++i)
  49. ans+=(it.s*mp[p][it.f+i]);
  50. for(auto &it:mp[q])
  51. mp[p][it.f]+=it.s;
  52. }
  53.  
  54. int main(){
  55. sp;
  56. init();
  57. int n,u,v;
  58. ll x;
  59. cin>>n>>d;
  60. for(int i=1; i<=n; ++i){
  61. cin>>x;
  62. mp[i][x]=1;
  63. }
  64. vector<pi> g;
  65. g.pb({0,0});
  66. for(int i=1; i<n; ++i){
  67. cin>>u>>v;
  68. g.pb({u,v});
  69. }
  70. int c[n];
  71. for(int i=1; i<n; ++i)
  72. cin>>c[i];
  73. ll b[n];
  74. for(int i=n-1; i>0; --i){
  75. b[i]=ans;
  76. unn(g[c[i]].f,g[c[i]].s);
  77. }
  78. for(int i=1; i<n; ++i)
  79. cout<<b[i]<<"\n";
  80.  
  81. return 0;}
Success #stdin #stdout 0.01s 8548KB
stdin
5 5
3 -2 -1 -4 0 
5 3
4 3
1 4
2 4
4
1
2
3
stdout
5
2
0
0