제출 #643288

#제출 시각아이디문제언어결과실행 시간메모리
643288valerikkConstruction of Highway (JOI18_construction)C++17
100 / 100
813 ms21880 KiB
#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int N = 1e5 + 7;

int n;
int c[N];
int a[N], b[N];
vector<int> g[N];
int sz[N];
vector<int> order;
int tin[N], tout[N];
int top[N];
int par[N];

void dfs(int u) {
  tin[u] = order.size();
  order.push_back(u);
  
  for (int &v : g[u]) {
    if (sz[v] > sz[g[u][0]]) 
      swap(v, g[u][0]);
  }

  for (int v : g[u]) {
    top[v] = v == g[u][0] ? top[u] : v;
    dfs(v);
  }

  tout[u] = order.size();
}

int mn[4 * N], mx[4 * N], color[4 * N];

void apply(int v, int c) {
  mn[v] = c;
  mx[v] = c;
  color[v] = c;
}

void push(int v) {
  if (color[v] == -1)
    return;

  apply(2 * v, color[v]);
  apply(2 * v + 1, color[v]);
  color[v] = -1;
}

void update(int v, int left, int right, int l, int r, int c) {
  if (left >= r || right <= l)
    return;
  if (left >= l && right <= r) {
    apply(v, c);
    return;
  }

  push(v);

  int mid = (left + right) / 2;

  update(2 * v, left, mid, l, r, c);
  update(2 * v + 1, mid, right, l, r, c);

  mn[v] = min(mn[2 * v], mn[2 * v + 1]);
  mx[v] = max(mx[2 * v], mx[2 * v + 1]);
}

int get_color(int v, int left, int right, int pos) {
  if (color[v] != -1) 
    return color[v];

  if (right - left == 1)
    return c[order[left]];

  int mid = (left + right) / 2;

  if (pos < mid) 
    return get_color(2 * v, left, mid, pos);
  return get_color(2 * v + 1, mid, right, pos);
}

int find_prev(int v, int left, int right, int l, int r, int c) {
  if (left >= r || right <= l || (mn[v] == c && mx[v] == c))
    return -1;
  if (right - left == 1)
    return left;
  
  push(v);

  int mid = (left + right) / 2;

  int res = find_prev(2 * v + 1, mid, right, l, r, c);
  if (res == -1)
    res = find_prev(2 * v, left, mid, l, r, c);

  return res;
}

void build(int v, int left, int right) {
  color[v] = -1;

  if (right - left == 1) {
    mn[v] = mx[v] = c[order[left]];
    return;
  }

  int mid = (left + right) / 2;
  
  build(2 * v, left, mid);
  build(2 * v + 1, mid, right);

  mn[v] = min(mn[2 * v], mn[2 * v + 1]);
  mx[v] = max(mx[2 * v], mx[2 * v + 1]);
}

int fen[N];

void add_fen(int i, int val) {
  for (++i; i < N; i += i & -i) {
    fen[i] += val;
  }
}

int get_fen(int i) {
  int sum = 0;
  for (; i > 0; i -= i & -i) {
    sum += fen[i];
  }
  return sum;
}

long long get_cost(int u) {
  vector<pair<int, int>> segments;

  while (u != -1) {
    int c = get_color(1, 0, n, tin[u]);

    int pos = find_prev(1, 0, n, 0, tin[u] + 1, c);

    if (pos < tin[top[u]]) {
      segments.push_back({c, tin[u] - tin[top[u]] + 1});
      u = par[top[u]];
    } else {
      segments.push_back({c, tin[u] - pos});
      u = order[pos];
    }
  }

  long long res = 0;

  for (auto segment : segments) {
    // cout << segment.first << " " << segment.second << endl;
    res += get_fen(segment.first) * 1LL * segment.second;
    add_fen(segment.first, segment.second);
  }
  for (auto segment : segments) {
    add_fen(segment.first, -segment.second);
  }

  return res; 
}

int main() { 
  cin >> n;
  for (int i = 0; i < n; ++i) {
    cin >> c[i];
  }

  vector<int> sorted;
  for (int i = 0; i < n; ++i) {
    sorted.push_back(c[i]);
  }
  sort(sorted.begin(), sorted.end());
  for (int i = 0; i < n; ++i) {
    c[i] = lower_bound(sorted.begin(), sorted.end(), c[i]) - sorted.begin();
  }

  for (int i = 0; i < n - 1; ++i) {
    cin >> a[i] >> b[i];
    --a[i];
    --b[i];
  }

  par[0] = -1;
  for (int i = 0; i < n - 1; ++i) {
      par[b[i]] = a[i];
      g[a[i]].push_back(b[i]); 
  }

  for (int i = 0; i < n; ++i) {
    sz[i] = 1;
  }
  for (int i = n - 1; i >= 0; --i) {
    sz[a[i]] += sz[b[i]];
  }

  dfs(0);

  build(1, 0, n);

  for (int i = 0; i < n - 1; ++i) {
    cout << get_cost(a[i]) << "\n";

    // for (int u = 0; u < n; ++u) {
    //   cout << get_color(1, 0, n, tin[u]) << " ";
    // }
    // cout << endl;

    for (int u = b[i]; u != -1; u = par[top[u]]) {
      update(1, 0, n, tin[top[u]], tin[u] + 1, c[b[i]]);
    }
  }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...