fork download
  1. #include <bits/stdc++.h>
  2. #define endl '\n'
  3.  
  4. using namespace std;
  5. const int MAXN = (1 << 20);
  6. const int MAXLOG = 20;
  7.  
  8. template <class T>
  9. struct fenwick
  10. {
  11. int sz;
  12. T tr[MAXN];
  13.  
  14. void init(int n)
  15. {
  16. sz = n + 1;
  17. memset(tr, 0, sizeof(tr));
  18. }
  19.  
  20. T query(int idx)
  21. {
  22. T ans = 0;
  23. for(; idx >= 1; idx -= (idx & -idx))
  24. ans += tr[idx];
  25. return ans;
  26. }
  27.  
  28. void update(int idx, T val)
  29. {
  30. if(idx <= 0) return;
  31. for(; idx <= sz; idx += (idx & -idx))
  32. tr[idx] += val;
  33. }
  34.  
  35. T query(int l, int r) { return query(r) - query(l - 1); }
  36. };
  37.  
  38. int n, h;
  39. vector<int> G[MAXN];
  40.  
  41. void read()
  42. {
  43. cin >> n >> h;
  44. for(int i = 1; i < n; i++)
  45. {
  46. int p;
  47. cin >> p;
  48. G[p].push_back(i);
  49. G[i].push_back(p);
  50. }
  51. }
  52.  
  53. int lgn, st[MAXN], en[MAXN], depth[MAXN], dfs_time = 0;
  54. int par[MAXN][MAXLOG], root = 0;
  55. vector<int> li[MAXN];
  56.  
  57. void dfs_lca(int u, int d = 0)
  58. {
  59. li[d].push_back(u);
  60. depth[u] = d;
  61.  
  62. int sz = G[u].size(), v;
  63. st[u] = ++dfs_time;
  64. for(int i = 1; i < MAXLOG; i++)
  65. par[u][i] = par[par[u][i - 1]][i - 1];
  66.  
  67. for(int i = 0; i < sz; i++)
  68. {
  69. v = G[u][i];
  70. if(v != par[u][0])
  71. {
  72. par[v][0] = u;
  73. dfs_lca(v, d + 1);
  74. }
  75. }
  76.  
  77. en[u] = dfs_time;
  78. }
  79.  
  80. inline bool upper(int u, int v) { return st[u] <= st[v] && en[v] <= en[u]; }
  81.  
  82. int lca(int u, int v)
  83. {
  84. if(upper(u, v)) return u;
  85. if(upper(v, u)) return v;
  86.  
  87. int a = u;
  88. for(int i = MAXLOG - 1; i >= 0; i--)
  89. if(!upper(par[a][i], v))
  90. a = par[a][i];
  91.  
  92. return par[a][0];
  93. }
  94.  
  95. int parent(int u, int up)
  96. {
  97. int a = u;
  98. for(int i = MAXLOG - 1; i >= 0; i--)
  99. if(up & (1 << i))
  100. a = par[a][i];
  101.  
  102. return a;
  103. }
  104.  
  105. void lca_precompute(int _root)
  106. {
  107. root = _root;
  108. for(int i = 0; i < MAXLOG; i++) par[root][i] = root;
  109. dfs_time = 0;
  110. dfs_lca(root);
  111. }
  112.  
  113. fenwick<int64_t> anst, t;
  114.  
  115. int get_par(int u, int k)
  116. {
  117. if(t.query(st[u], en[u]) >= k) return u;
  118. for(int l = MAXLOG - 1; l >= 0; l--)
  119. if(t.query(st[par[u][l]], en[par[u][l]]) < k)
  120. u = par[u][l];
  121.  
  122. u = par[u][0];
  123. if(t.query(st[u], en[u]) < k) return -1;
  124. return u;
  125. }
  126.  
  127. int get_real_par(int u)
  128. {
  129. for(int l = MAXLOG - 1; l >= 0; l--)
  130. if(!anst.query(st[par[u][l]], en[par[u][l]]))
  131. u = par[u][l];
  132.  
  133. return u;
  134. }
  135.  
  136. int64_t solve(int l, int k)
  137. {
  138. int64_t ret = 0;
  139. vector<int> to_rem;
  140.  
  141. for(int u: li[l])
  142. {
  143. int up, p = get_par(u, k);
  144. if(p == -1) continue;
  145. if(anst.query(st[p], en[p]))
  146. continue;
  147.  
  148. up = get_real_par(p);
  149. ret += depth[p] - depth[up] + 1;
  150. anst.update(st[p], 1);
  151. to_rem.push_back(p);
  152. }
  153.  
  154. for(int u: to_rem)
  155. anst.update(st[u], -1);
  156.  
  157. return ret;
  158. }
  159.  
  160. void solve()
  161. {
  162. lca_precompute(0);
  163. anst.init(n + 2);
  164. t.init(n + 2);
  165.  
  166. int64_t ans = 0;
  167. for(int i = 0; i <= h; i++)
  168. {
  169. for(int u: li[i]) t.update(st[u], 1);
  170.  
  171. int k;
  172. cin >> k;
  173.  
  174. if(k != 0) ans += 1ll * solve(i, k);
  175. else ans += 1ll * n;
  176.  
  177. for(int u: li[i]) t.update(st[u], -1);
  178. }
  179.  
  180. cout << ans << endl;
  181. }
  182.  
  183. int main()
  184. {
  185. ios_base::sync_with_stdio(false);
  186. cin.tie(NULL);
  187.  
  188. read();
  189. solve();
  190. return 0;
  191. }
  192.  
Success #stdin #stdout 0.01s 175808KB
stdin
Standard input is empty
stdout
0