제출 #1059220

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1059220		MilosMilutinovic	Unique Cities (JOI19_ho_t5)	C++14	4 / 100	2049 ms	274432 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, m;
  cin >> n >> m;
  vector<vector<int>> g(n);
  for (int i = 1; i < n; i++) {
    int x, y;
    cin >> x >> y;
    --x; --y;
    g[x].push_back(y);
    g[y].push_back(x);
  }
  vector<int> c(n);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
    --c[i];
  }
  auto Bfs = [&](int v) {
    vector<int> d(n, -1);
    vector<int> que(1, v);
    d[v] = 0;
    for (int b = 0; b < (int) que.size(); b++) {
      int i = que[b];
      for (int j : g[i]) {
        if (d[j] == -1) {
          d[j] = d[i] + 1;
          que.push_back(j);
        }
      }
    }
    return d;
  }; 
  int x, y;
  {
    vector<int> d = Bfs(0);
    x = 0;
    for (int i = 1; i < n; i++) {
      if (d[i] > d[x]) {
        x = i;
      }
    }
  }
  {
    vector<int> d = Bfs(x);
    y = 0;
    for (int i = 1; i < n; i++) {
      if (d[i] > d[y]) {
        y = i;
      }
    }
  }
  vector<vector<int>> d(2);
  d[0] = Bfs(x);
  d[1] = Bfs(y);
  vector<int> res(n);
  for (int foo = 0; foo < 2; foo++) {
    int root = (foo == 0 ? x : y);
    const int L = 18;
    vector<int> mx(n);
    vector<vector<int>> pr(n, vector<int>(L));
    vector<int> dep(n);
    function<void(int, int)> Dfs = [&](int v, int pv) {
      pr[v][0] = pv;
      for (int i = 1; i < L; i++) {
        pr[v][i] = pr[pr[v][i - 1]][i - 1];
      }
      mx[v] = 1;
      for (int u : g[v]) {
        if (u == pv) {
          continue;
        }
        dep[u] = dep[v] + 1;
        Dfs(u, v);
        mx[v] = max(mx[v], mx[u] + 1);
      }
    };
    auto Jump = [&](int v, int k) {
      for (int i = L - 1; i >= 0; i--) {
        if (k >> i & 1) {
          v = pr[v][i];
        }
      }
      return v;
    };
    Dfs(root, root);
    vector<int> st(8 * n, 1);
    vector<int> lzy(8 * n);
    auto Push = [&](int x) {
      st[x * 2] += lzy[x];
      st[x * 2 + 1] += lzy[x];
      lzy[x * 2] += lzy[x];
      lzy[x * 2 + 1] += lzy[x];
      lzy[x] = 0;
    }; 
    function<void(int, int, int, int, int, int)> Modify = [&](int x, int l, int r, int ll, int rr, int v) {
      if (ll <= l && r <= rr) {
        st[x] += v;
        lzy[x] += v;
        Push(x);
        return;
      }
      Push(x);
      int mid = (l + r) >> 1;
      if (rr <= mid) {
        Modify(x * 2, l, mid, ll, rr, v);
      } else if (ll > mid) {
        Modify(x * 2 + 1, mid + 1, r, ll, rr, v);
      } else {
        Modify(x * 2, l, mid, ll, rr, v);
        Modify(x * 2 + 1, mid + 1, r, ll, rr, v);
      }
      st[x] = max(st[x * 2], st[x * 2 + 1]);
    };
    function<int(int, int, int, int, int)> Query = [&](int x, int l, int r, int ll, int rr) {
      if (ll <= l && r <= rr) {
        return st[x];
      }
      Push(x);
      int mid = (l + r) >> 1;
      int res;
      if (rr <= mid) {
        res = Query(x * 2, l, mid, ll, rr);
      } else if (ll > mid) {
        res = Query(x * 2 + 1, mid + 1, r, ll, rr);
      } else {
        res = max(Query(x * 2, l, mid, ll, rr), Query(x * 2 + 1, mid + 1, r, ll, rr));
      }
      st[x] = max(st[x * 2], st[x * 2 + 1]);
      return res;
    };
    auto Add = [&](int v, int d, int p) {
      Modify(1, 0, n - 1, max(0, dep[v] - d + 1), dep[v], p);
    };
    vector<set<int>> cols(n);
    function<void(int, int)> Solve = [&](int v, int pv) {
      if (d[foo][v] >= d[foo ^ 1][v]) {
        Add(v, mx[v], -1);
        vector<bool> was(m);
        int low = 0, high = dep[v] - 1, d = -1;
        while (low <= high) {
          int mid = (low + high) >> 1;
          if (Query(1, 0, n - 1, mid, dep[v] - 1) == 1) {
            d = mid;
            low = mid + 1;
          } else {
            high = mid - 1;
          }
        }
        if (d == -1) {
          res[v] = 0;
        } else {
          res[v] = (int) cols[Jump(v, dep[v] - d)].size();
        }
        Add(v, mx[v], +1);
      }
      multiset<int> st;
      for (int u : g[v]) {
        if (u == pv) {
          continue;
        }
        st.insert(mx[u]);
      }
      for (int u : g[v]) {
        if (u == pv) {
          continue;
        }
        st.erase(st.find(mx[u]));
        if (!st.empty()) {
          Add(pr[v][0], *prev(st.end()), -1);
        }
        cols[v].clear();
        int low = 0, high = dep[v] - 1, d = -1;
        while (low <= high) {
          int mid = (low + high) >> 1;
          if (Query(1, 0, n - 1, mid, dep[v] - 1) == 1) {
            d = mid;
            low = mid + 1; 
          } else {
            high = mid - 1;
          }
        }
        if (d != -1) {
          cols[v] = cols[Jump(v, dep[v] - d)];
        }
        cols[v].insert(c[v]);
        Solve(u, v);
        if (!st.empty()) {
          Add(pr[v][0], *prev(st.end()), +1);
        }
        st.insert(mx[u]);
      }
    };
    Solve(root, root);
  }
  for (int i = 0; i < n; i++) {
    cout << res[i] << '\n';
  }
  return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 0 / 32
• Subtask 3: 0 / 32
• Subtask 4: 0 / 32