Submission #1059244

#	Time	Username	Problem	Language	Result	Execution time	Memory
1059244		MilosMilutinovic	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	1282 ms	154008 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

const int MAX = (int) (8e6);

int rt[MAX], ch[MAX][2], cnt[MAX], tsz;

void Change(int& x, int y, int l, int r, int p) {
  x = ++tsz;
  ch[x][0] = ch[y][0];
  ch[x][1] = ch[y][1];
  cnt[x] = cnt[y];
  if (l == r) {
    cnt[x] = 1;
    return;
  }
  int mid = (l + r) >> 1;
  if (p <= mid) {
    Change(ch[x][0], ch[y][0], l, mid, p);
  } else {
    Change(ch[x][1], ch[y][1], mid + 1, r, p);
  }
  cnt[x] = cnt[ch[x][0]] + cnt[ch[x][1]]; 
}

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;
    };
    function<int(int, int, int, int)> Walk = [&](int x, int l, int r, int p) {
      if (l > p || st[x] < 1) {
        return -1;
      }
      if (l == r) {
        return l;
      }
      Push(x);
      int mid = (l + r) >> 1;
      int ret = Walk(x * 2 + 1, mid + 1, r, p);
      if (ret == -1) {
        ret = Walk(x * 2, l, mid, p);
      }
      st[x] = max(st[x * 2], st[x * 2 + 1]);
      return ret;
    };
    auto Add = [&](int v, int d, int p) {
      Modify(1, 0, n - 1, max(0, dep[v] - d + 1), dep[v], p);
    };
    while (tsz > 0) {
      ch[tsz][0] = 0;
      ch[tsz][1] = 0;
      cnt[tsz] = 0;
      tsz--;
    }
    for (int i = 0; i < n; i++) {
      rt[i] = 0;
    }
    function<void(int, int)> Solve = [&](int v, int pv) {
      if (d[foo][v] >= d[foo ^ 1][v]) {
        Add(v, mx[v], -1);
        int d = (dep[v] == 0 ? -1 : Walk(1, 0, n - 1, dep[v] - 1));
        if (d == -1) {
          res[v] = 0;
        } else {
          res[v] = cnt[rt[Jump(v, dep[v] - d)]];
        }
        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);
        }
        int d = (dep[v] == 0 ? -1 : Walk(1, 0, n - 1, dep[v] - 1));
        rt[v] = 0;
        if (d != -1) {
          rt[v] = rt[Jump(v, dep[v] - d)];
        }
        Change(rt[v], rt[v], 0, m - 1, 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: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32