제출 #220462

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
220462		rama_pang	Construction of Highway (JOI18_construction)	C++14	100 / 100	407 ms	108892 KiB

construction

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100005;

int N;
int A[MAXN], B[MAXN], C[MAXN];
vector<int> adj[MAXN];

void Read() {
  ios::sync_with_stdio(0);
  cin.tie(0), cout.tie(0);

  cin >> N;
  
  for (int i = 1; i <= N; i++) {
    cin >> C[i];
  }

  for (int i = 0; i < N - 1; i++) {
    cin >> A[i] >> B[i];
    adj[A[i]].emplace_back(B[i]);
    adj[B[i]].emplace_back(A[i]);
  }

  // Coordinate compression on C
  vector<int> coord;
  for (int i = 1; i <= N; i++) {
    coord.emplace_back(C[i]);
  }
  sort(begin(coord), end(coord));
  coord.resize(unique(begin(coord), end(coord)) - begin(coord));
  for (int i = 1; i <= N; i++) {
    C[i] = lower_bound(begin(coord), end(coord), C[i]) - begin(coord) + 1;
  }
}

//// Heavy Light Decomposition
int root[MAXN]; // root of current chain
int par[MAXN]; // parent node
int sz[MAXN]; // size of subtree
int depth[MAXN];

void dfsSz(int n, int p) {
  if (p) adj[n].erase(find(begin(adj[n]), end(adj[n]), p));
  par[n] = p;
  sz[n] = 1;
  depth[n] = depth[p] + 1;
  for (auto &i : adj[n]) {
    dfsSz(i, n);
    sz[n] += sz[i];
    if (sz[i] > sz[adj[n][0]]) {
      swap(i, adj[n][0]);
    }
  }
}

void dfsHld(int n) {
  for (auto &i : adj[n]) {
    root[i] = (i == adj[n][0] ? root[n] : i);
    dfsHld(i);
  }
}

void HLD() {
  root[1] = 1;
  dfsSz(1, 0);
  dfsHld(1);
}

//// Binary Indexed Tree
int BIT[MAXN];

void Add(int pos, int x) {
  for (int i = pos; i <= N; i += i & -i) BIT[i] += x;
}

int Sum(int pos) {
  int res = 0;
  for (int i = pos; i > 0; i -= i & -i) res += BIT[i];
  return res;
}

//// Main Solution
long long CountInversion(const vector<pair<int, int>> &V) { // (frequency, value)
  long long res = 0;
  int n = V.size();
  for (int i = n - 1; i >= 0; i--) {
    Add(V[i].second, V[i].first);
    res += 1ll * V[i].first * Sum(V[i].second - 1);
  }
  for (int i = 0; i < n; i++) {
    Add(V[i].second, -V[i].first);
  }
  
  return res;
}

deque<pair<int, int>> chainValues[MAXN]; // values in the chain, (frequency, value). Each construction adds at most O(log N) values.

long long Construction(int a, int b) { // construct a new edge (a, b)
  vector<pair<int, int>> path; // values in path from root to a

  // get values in path and delete them
  while (a != 0) {
    deque<pair<int, int>> &chain = chainValues[root[a]];
    int take = depth[a] - depth[root[a]] + 1;
    int ptr = path.size();
    while (take > 0) {
      if (take >= chain.front().first) {
        take -= chain.front().first;
        path.emplace_back(chain.front());
        chain.pop_front();
      } else {
        path.emplace_back(take, chain.front().second);
        chain.front().first -= take;
        take = 0;
      }
    }
    reverse(begin(path) + ptr, end(path));
    a = par[root[a]];
  }
  
  reverse(begin(path), end(path));

  int new_C = C[b];
  while (b != 0) {
    deque<pair<int, int>> &chain = chainValues[root[b]];
    int put = depth[b] - depth[root[b]] + 1;
    chain.emplace_front(put, new_C);
    b = par[root[b]];
  }

  return CountInversion(path);
}

void Solve() {
  chainValues[1].emplace_back(1, C[1]);
  for (int i = 0; i < N - 1; i++) {
    cout << Construction(A[i], B[i]) << "\n";
  }
}

int main() {
  Read();
  HLD();
  Solve();
  return 0;
}

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