Submission #624602

#	Time	Username	Problem	Language	Result	Execution time	Memory
624602		MilosMilutinovic	Construction of Highway (JOI18_construction)	C++14	100 / 100	1004 ms	34712 KiB

construction

/**
 *    author:  wxhtzdy
 *    created: 08.08.2022 14:06:23
**/
#include <bits/stdc++.h>

using namespace std;

template <typename T>
class fenwick {
  public:
  vector<T> fenw;
  int n;
  fenwick(int _n) : n(_n) {
    fenw.resize(n);
  }
  void modify(int x, T v) {
    while (x < n) {
      fenw[x] += v;
      x |= (x + 1);
    }
  }
  T get(int x) {
    T v{};
    while (x >= 0) {
      v += fenw[x];
      x = (x & (x + 1)) - 1;
    }
    return v;
  }
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);  
  int n;
  cin >> n;
  vector<int> c(n);
  for (int i = 0; i < n; i++) {
    cin >> c[i];
  }
  auto xs = c;
  sort(xs.begin(), xs.end());
  xs.erase(unique(xs.begin(), xs.end()), xs.end());
  for (int i = 0; i < n; i++) {
    c[i] = lower_bound(xs.begin(), xs.end(), c[i]) - xs.begin();
  }  
  vector<vector<int>> g(n);
  vector<pair<int, int>> edges;
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    --u; --v;                
    g[u].push_back(v);
    g[v].push_back(u);
    edges.emplace_back(u, v);
  }
  vector<int> par(n);
  vector<int> dep(n);
  vector<int> tin(n);
  vector<int> tout(n);
  int T = 0;   
  function<void(int, int)> Dfs = [&](int v, int pr) {
    tin[v] = ++T;
    par[v] = pr;
    dep[v] = dep[pr] + 1;
    for (int u : g[v]) {
      if (u != pr) {
        Dfs(u, v);
      }
    }
    tout[v] = T;
  };
  Dfs(0, 0);
  const int L = 25;
  vector<vector<int>> jump(n, vector<int>(L));
  for (int i = 0; i < n; i++) {
    jump[i][0] = par[i];  
  }
  for (int j = 1; j < L; j++) {
    for (int i = 0; i < n; i++) {
      jump[i][j] = jump[jump[i][j - 1]][j - 1];
    }
  }     
  auto isPar = [&](int u, int v) {
    return tin[u] <= tin[v] && tout[v] <= tout[u];
  };
  auto Find = [&](int u, int v) {
    for (int i = L - 1; i >= 0; i--) {
      if (!isPar(jump[v][i], u)) {
        v = jump[v][i];
      }
    }
    return v;
  };                   
  vector<int> f(n, -1);
  f[0] = 0;
  fenwick<int> fenw(n);
  for (auto& p : edges) {
    int a = p.first;
    int b = p.second;    
    int x = 0;
    vector<pair<int, int>> v;
    while (x != b && !isPar(f[x], b)) {
      int i = Find(f[x], b);
      int j = Find(b, f[x]);
      f[j] = f[x];
      v.emplace_back(dep[i] - dep[x], c[f[x]]);        
      x = i;
    }   
    if (x != b) {      
      v.emplace_back(dep[b] - dep[x], c[f[x]]);
    }
    f[0] = b;
    long long ans = 0;
    reverse(v.begin(), v.end());
    for (auto& p : v) {
      ans += fenw.get(p.second - 1) * 1LL * p.first;
      fenw.modify(p.second, p.first);
    }
    cout << ans << '\n';
    for (auto& p : v) {
      fenw.modify(p.second, -p.first);    
    }
  } 
  return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:100:9: warning: unused variable 'a' [-Wunused-variable]
  100 |     int a = p.first;
      |         ^

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