Submission #1035358

#	Time	Username	Problem	Language	Result	Execution time	Memory
1035358		juicy	Construction of Highway (JOI18_construction)	C++17	100 / 100	429 ms	28856 KiB

construction

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif

const int N = 1e5 + 5;

int n;
int c[N], pr[N], head[N], tin[N], sz[N], s[N], bot[N];
set<array<int, 2>> st[N];
vector<int> g[N];

void upd(int i, int x) {
  for (; i <= n; i += i & -i) {
    s[i] += x;
  }
}

int qry(int i) {
  int res = 0;
  for (; i; i -= i & -i) {
    res += s[i];
  }
  return res;
}

void dfs(int u) {
  for (int &v : g[u]) {
    if (v == pr[u]) {
      swap(v, g[u].back());
      g[u].pop_back();
      break;
    }
  }
  sz[u] = 1;
  for (int v : g[u]) {
    pr[v] = u;
    dfs(v);
    sz[u] += sz[v];
  }
}

int timer;

void hld(int u, int h) {
  tin[u] = ++timer;
  head[u] = h;
  bot[h] = timer;
  if (g[u].size()) {
    int x = *max_element(g[u].begin(), g[u].end(), [&](int i, int j) {
      return sz[i] < sz[j];
    });
    hld(x, h);
    for (int v : g[u]) {
      if (v != x) {
        hld(v, v);
      }
    }
  }
}

long long qry(int u, int x) {
  vector<array<int, 2>> cands;
  long long res = 0;
  for (; u; u = pr[head[u]]) {
    int h = head[u];
    auto it = st[h].lower_bound({tin[u] + 1, 0});
    assert(it != st[h].end());
    bool exist = (*it)[0] == tin[u] + 1; 
    --it;
    if (!exist) {
      st[h].insert({tin[u] + 1, (*it)[1]});
    }
    vector<array<int, 2>> nw;
    while (st[h].size() && (*(it = st[h].begin()))[0] <= tin[u]) {
      nw.push_back({(*next(it))[0] - (*it)[0], (*it)[1]});
      st[h].erase(it);
    }
    reverse(nw.begin(), nw.end());
    cands.insert(cands.end(), nw.begin(), nw.end());
    nw.clear();
    st[h].insert({tin[h], x});
  }
  for (auto [l, x] : cands) {
    res += (long long) qry(x - 1) * l;
    upd(x, l);
  }
  for (auto [l, x] : cands) {
    upd(x, -l);
  }
  cands.clear();
  return res;
}

int main() {
  ios::sync_with_stdio(false); cin.tie(nullptr);

  cin >> n;
  vector<int> comp;
  for (int i = 1; i <= n; ++i) {
    cin >> c[i];
    comp.push_back(c[i]);
  }
  vector<array<int, 2>> edges;
  for (int i = 1; i < n; ++i) {
    int u, v; cin >> u >> v;
    edges.push_back({u, v});
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1);
  hld(1, 1);
  sort(comp.begin(), comp.end());
  comp.erase(unique(comp.begin(), comp.end()), comp.end());
  for (int i = 1; i <= n; ++i) {
    int h = head[i];
    c[i] = lower_bound(comp.begin(), comp.end(), c[i]) - comp.begin() + 1;
    st[h].insert({tin[i], c[i]});
    if (bot[i]) {
      st[h].insert({bot[i] + 1, -1});
    }
  }
  for (auto [u, v] : edges) {
    cout << qry(u, c[v]) << "\n";
  }
  return 0;
}

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84