Submission #1349059

#TimeUsernameProblemLanguageResultExecution timeMemory
1349059avighnaCollecting Stamps 5 (JOI26_stamps)C++20
100 / 100
1125 ms55336 KiB
#include <bits/stdc++.h>

using namespace std;

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);

  int n, d;
  cin >> n >> d;
  vector<int> t(n);
  for (int &i : t) {
    cin >> i;
  }
  vector<vector<int>> adj(n);
  for (int i = 0, u, v; i < n - 1; ++i) {
    cin >> u >> v;
    --u, --v;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }

  vector<bool> is_alr_centroid(n);

  vector<int> subtree_sz(n);
  auto find_subtree_sz = [&](auto &&self, int u, int p) -> void {
    subtree_sz[u] = 1;
    for (int &i : adj[u]) {
      if (i != p && !is_alr_centroid[i]) {
        self(self, i, u);
        subtree_sz[u] += subtree_sz[i];
      }
    }
  };

  auto find_centroid = [&](auto &&self, int u, int p, int n) {
    pair<int, int> mx = {-1, -1};
    for (int &i : adj[u]) {
      if (i != p && !is_alr_centroid[i]) {
        mx = max(mx, {subtree_sz[i], i});
      }
    }
    if (mx.first <= n / 2) {
      return u;
    }
    return self(self, mx.second, u, n);
  };

  vector<int> dep(n), ans(n), lb(n);
  vector<bool> good_up(n);
  auto process = [&](int c) {
    int max_dep = 0;
    vector<int> nodes;
    auto dfs = [&](auto &&self, int u, int p, int val1, int val2) -> void {
      nodes.push_back(u);
      max_dep = max(max_dep, dep[u]);
      good_up[u] = dep[u] >= val1;
      lb[u] = val2;
      for (int &i : adj[u]) {
        if (i == p || is_alr_centroid[i]) {
          continue;
        }
        dep[i] = dep[u] + 1;
        self(self, i, u, min(val1, t[i] + dep[i]), min(val2, t[i] - dep[i]));
      }
    };
    dep[c] = 0;
    dfs(dfs, c, -1, t[c] + dep[c], t[c] - dep[c]);
    vector<int> atdep(max_dep + 1);
    for (int &i : nodes) {
      atdep[dep[i]]++;
      lb[i] = max(lb[i], 0);
    }
    partial_sum(atdep.begin(), atdep.end(), atdep.begin()); 
    auto at_dep = [&](int i) { return i < atdep.size() ? atdep[i] : atdep.back(); };
    if (t[c] == 0) {
      ans[c] += at_dep(d);
    }
    auto dfs2 = [&](auto &&self, int u, int p) -> void {
      nodes.push_back(u);
      max_dep = max(max_dep, dep[u]);
      for (int &i : adj[u]) {
        if (i != p && !is_alr_centroid[i]) {
          self(self, i, u);
        }
      }
    };
    vector<int> diff(max_dep + 2);
    auto all_nodes = nodes;
    for (int &subt : adj[c]) {
      if (is_alr_centroid[subt]) {
        continue;
      }
      nodes.clear();
      max_dep = 0;
      dfs2(dfs2, subt, c);
      vector<int> atdep_cur(max_dep + 1), diff_cur(max_dep + 2);
      for (int &i : nodes) {
        atdep_cur[dep[i]]++;
      }
      partial_sum(atdep_cur.begin(), atdep_cur.end(), atdep_cur.begin());
      auto at_dep_cur = [&](int i) { return i < atdep_cur.size() ? atdep_cur[i] : atdep_cur.back(); };
      auto incr = [&](int l, int r, int del) {
        if (l < diff.size()) {
          diff[l] += del;
        }
        if (r + 1 < diff.size()) {
          diff[r + 1] -= del;
        }
      };
      auto incr_cur = [&](int l, int r, int del) {
        if (l < diff_cur.size()) {
          diff_cur[l] += del;
        }
        if (r + 1 < diff_cur.size()) {
          diff_cur[r + 1] -= del;
        }
      };
      for (int &i : nodes) {
        if (good_up[i]) {
          if (d - dep[i] >= 0) {
            ans[i] += at_dep(d - dep[i]) - at_dep_cur(d - dep[i]);
          }
        }
        if (lb[i] <= d - dep[i]) {
          incr(lb[i], d - dep[i], 1);
          incr_cur(lb[i], d - dep[i], -1);
        }
      }
      partial_sum(diff_cur.begin(), diff_cur.end(), diff_cur.begin());
      for (int &i : nodes) {
        if (!good_up[i]) {
          ans[i] += diff_cur[dep[i]];
        }
      }
    }
    partial_sum(diff.begin(), diff.end(), diff.begin());
    for (int &i : all_nodes) {
      if (!good_up[i]) {
        ans[i] += diff[dep[i]];
      }
    }
  };

  auto centroid_decomp = [&](auto &&self, int u) -> void {
    find_subtree_sz(find_subtree_sz, u, -1);
    int c = find_centroid(find_centroid, u, -1, subtree_sz[u]);
    process(c);
    is_alr_centroid[c] = true;
    for (int &i : adj[c]) {
      if (!is_alr_centroid[i]) {
        self(self, i);
      }
    }
  };
  centroid_decomp(centroid_decomp, 0);

  for (int &i : ans) {
    cout << i << '\n';
  }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...