Submission #999836

#	Time	Username	Problem	Language	Result	Execution time	Memory
999836		vjudge1	Construction of Highway (JOI18_construction)	C++17	100 / 100	1736 ms	36024 KiB

construction

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#pragma GCC optimize("O3") 
using namespace __gnu_pbds;
using namespace std;
template <typename T>
using ordered_set =
    tree<T, null_type, less_equal<T>, rb_tree_tag, tree_order_statistics_node_update>;

const int N = 100'030;
const int LG = 20;
int up[N][LG], depth[N], n, sz[N], tin[N], tout[N], tim=0;
vector<int> g[N];
int c[N];
void dfs(int at, int p) {
  up[at][0] = p;
  sz[at]=1;
  for(int i = 1;i < LG;i++) {
    up[at][i] = up[up[at][i-1]][i-1];
  }
  for(int to:g[at]) {
    if(to == p)continue;
    depth[to] = depth[at]+1;
    dfs(to, at);
    sz[at]+=sz[to];
  }
}
int lca(int a, int b) {
  if(depth[a] < depth[b])swap(a, b);
  int k = depth[a] - depth[b];
  for(int i = LG-1;i>=0;i--) {
    if(k&(1<<i)) {
      a = up[a][i];
    }
  }
  if(a==b)return a;
  for(int i = LG-1;i>=0;i--) {
    if(up[a][i] != up[b][i]) {
      a=up[a][i];
      b=up[b][i];
    }
  }
  return up[a][0];
}
struct node {
  int mx, mn;
  node() {
    mx=-1e9;
    mn=1e9;
  }
};
node merge(node A, node B) {
  node res;
  res.mx = max(A.mx, B.mx);
  res.mn = min(A.mn, B.mn);
  return res;
}
struct DS1 {
  vector<node> T;
  vector<int> lz;
  DS1() {
    T.resize(4*N);
    lz=vector<int>(4*N, -1);
  }
//  void update(int l, int r, int val) {
//    for(int i = l;i<=r;i++)T[i].mx = T[i].mn = val;
//  }
//  node query(int l, int r) {
//    node s;
//    for(int i = l;i<=r;i++)s=merge(s, T[i]);
//    return s;
//  }
  void push(int v) {
    if(lz[v] == -1)return;
    lz[2*v] = lz[v];
    lz[2*v+1] = lz[v];
    T[2*v].mx = T[2*v].mn = lz[v];
    T[2*v+1].mx = T[2*v+1].mn = lz[v];
    lz[v] = -1;
  }
  void update(int l, int r, int val, int tl=0, int tr=N-1, int v=1) {
    if(l>tr or r<tl)return;
    if(l <= tl and tr <= r) {
      lz[v] = val;
      T[v].mx = val;
      T[v].mn = val;
      return;
    }
    int tm=(tl+tr)/2;
    push(v);
    update(l, r, val, tl, tm, 2*v);
    update(l, r, val, tm+1, tr, 2*v+1);
    T[v] = merge(T[2*v], T[2*v+1]);
  }
  node query(int l, int r, int tl=0, int tr = N-1, int v=1) {
    if(l > tr or r<tl)return node();
    if(l <= tl and tr <= r)
      return T[v];
    int tm=(tl+tr)/2;
    push(v);
    return merge(query(l, r, tl, tm, 2*v), query(l, r, tm+1, tr, 2*v+1));
  }
};
namespace HLD {
  DS1 ds;
  int tin[N], top[N], tim=0, par[N], V[N];
  void dfs(int at, int p, int tp) {
    tin[at]=++tim;
    top[at]=tp;
    par[at]=p;
    int big=-1;
    for(int to:g[at]) {
      if(to==p)continue;
      if(big==-1 or sz[to]>sz[big]) {
        big=to;
      }
    }
    if(big==-1)return;
    dfs(big, at, tp);
    for(int to:g[at]) {
      if(to==big or to==p)continue;
      dfs(to, at, to);
    }
  }
  int upd(int u,int v, long long val) {
    int ans=0;
    while(top[u]!=top[v]) {
     if(depth[top[u]] < depth[top[v]])swap(u,v);
//     cout << top[u] << " " << u << " " << tin[top[u]] << " " << tin[u] << " " << val << endl;
     ds.update(tin[top[u]], tin[u], val);
     u=par[top[u]];
    }
    if(depth[u]>depth[v])swap(u,v);
    ds.update(tin[u], tin[v], val);
    return ans;
  }
  node query(int u,int v) {
    node ans;
    while(top[u]!=top[v]) {
     if(depth[top[u]] < depth[top[v]])swap(u,v);
     ans=merge(ans, ds.query(tin[top[u]], tin[u]));
     u=par[top[u]];
    }
    if(depth[u]>depth[v])swap(u,v);
    ans=merge(ans, ds.query(tin[u], tin[v]));
    return ans;
  }
  long long get(int u) {
    return ds.query(tin[u], tin[u]).mx;
  }
};
int get(int par) {
  for(int i = LG-1;i>=0;i--) {
    if(up[par][i]==0)continue;
    node res = HLD::query(up[par][i], par);
    if(res.mx == res.mn) {
      par = up[par][i];
    }
  }
  return up[par][0];
}
long long fen[N];
void upd(int at, int val) {
  at++;
  while(at<N) {
    fen[at]+=val;
    at+=at&(-at);
  }
}
long long que(int at) {
  if(at==0)return 0;
  long long res=0;
  while(at > 0) {
    res+=fen[at];
    at-=at&(-at);
  }
  return res;
}
int main () {
  cin.tie(0)->sync_with_stdio(0);
  cin >> n;
  for(int i = 1;i<=n;i++){
    cin >> c[i];
  }
  vector<pair<int,int>> e;
  for(int i = 1;i<n;i++) {
    int a, b;
    cin >> a >> b;
    g[a].push_back(b);
    g[b].push_back(a);
    e.push_back({a, b});
  }
  {
    map<int,int> mp;
    for(int i = 1;i<=n;i++)mp[c[i]] = 1;
    int cc = 0;
    for(auto &i:mp) {
      i.second=cc++;
    }
    for(int i = 1;i<=n;i++)c[i] = mp[c[i]];
  }
  depth[0]=-1;
  dfs(1, 0);
  HLD::dfs(1, 0, 1);
  int cnt=0;
  for(int i = 1;i<=n;i++) {
    HLD::upd(i, i, c[i]);
  }
  for(auto [a, b]:e) {
    long long ans = 0;
    int cur=a;
    int prev =0;
    vector<pair<int,int>> change;
    do {
      int par = get(cur);
      int val = HLD::get(cur);
      ans+=((long long)depth[cur]-depth[par])*que(val);
      upd(val, depth[cur]-depth[par]);
      change.push_back({val, depth[cur]-depth[par]});
      cur = par;
    }while(cur!=0);
    for(auto i:change) {
      upd(i.first, -i.second);
    }
    cur=a;
    HLD::upd(1, a, c[b]);
    cout << ans << "\n";
  }
  return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:214:9: warning: unused variable 'prev' [-Wunused-variable]
  214 |     int prev =0;
      |         ^~~~
construction.cpp:207:7: warning: unused variable 'cnt' [-Wunused-variable]
  207 |   int cnt=0;
      |       ^~~

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