Submission #659741

#	Time	Username	Problem	Language	Result	Execution time	Memory
659741		MohammadAghil	Construction of Highway (JOI18_construction)	C++17	100 / 100	363 ms	30160 KiB

construction

      #include <bits/stdc++.h>
//   #pragma GCC optimize ("Ofast,unroll-loops")
// #pragma GCC target ("avx2")
 using namespace std;
  typedef long long ll;
   typedef pair<int, int> pp;
     #define per(i,r,l) for(int i = (r); i >= (l); i--)
       #define rep(i,l,r) for(int i = (l); i < (r); i++)
          #define all(x) begin(x), end(x)
             #define sz(x) (int)(x).size()
                 #define pb push_back
                     #define ss second
                          #define ff first
                                  void err(istringstream *iss){}template<typename T,typename ...Args> void err(istringstream *iss,const T &_val, const Args&...args){string _name;*iss>>_name;if(_name.back()==',')_name.pop_back();cerr<<_name<<" = "<<_val<<", ",err(iss,args...);}
void IOS(){
     cin.tie(0) -> sync_with_stdio(0);
     // #ifndef ONLINE_JUDGE
     //      #define er(args ...) cerr << __LINE__ << ": ", err(new istringstream(string(#args)), args), cerr << endl
     //      freopen("inp.txt", "r", stdin);
     //      freopen("out.txt", "w", stdout);
     // #else
     //      #define er(args ...) 0
     // #endif
}
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
const ll mod = 1e9 + 7, maxn = 1e5 + 5, lg = 22, inf = ll(1e9) + 5;
ll pw(ll a,ll b,ll md=mod){if(!b)return 1;ll k=pw(a,b>>1ll,md);return k*k%md*(b&1ll?a:1)%md;}

int n, c[maxn], h[maxn], par[maxn], cnt[maxn], srt[maxn], top[maxn];
vector<int> adj[maxn];
priority_queue<pp, vector<pp>, greater<pp>> pq[maxn];

void dfs(int r, int p){
     cnt[r] = 1;
     par[r] = p;
     for(int c: adj[r]) if(c - p){
          h[c] = h[r] + 1;
          dfs(c, r);
          cnt[r] += cnt[c];
     }
}

void dfsh(int r, int p, int tp){
     top[r] = tp;
     // er(r, tp, srt[r]);
     pp bst = {-1, -1};
     pq[tp].push({srt[r], c[r]});
     for(int c: adj[r]) if(c - p) bst = max(bst, {cnt[c], c});
     if(bst.ss + 1){
          srt[bst.ss] = srt[r] + 1;
          dfsh(bst.ss, r, tp);
          for(int c: adj[r]) if(c - p && c - bst.ss) dfsh(c, r, c);
     }
}

ll fen[maxn];
vector<pp> his;
void upd(int i, int k, bool f = true){
     if(f) his.pb({i, k});
     for(; i < maxn; i += i&-i) fen[i] += k;
}
ll get(int i){
     ll res = 0;
     for(; i; i -= i&-i) res += fen[i];
     return res;
}
void clear(){
     for(auto[i, k]: his) upd(i, -k, false);
     his.clear();
} 

ll query(int u, int nw){
     ll res = 0;
     while(u + 1){
          vector<pp> tmp;
          int lst = -1;
          while(true){
               auto[tp, vl] = pq[top[u]].top();
               if(tp > srt[u]){
                    tmp.pb({srt[u] - lst, vl});
                    break;
               }else{
                    pq[top[u]].pop();
                    tmp.pb({tp - lst, vl});
                    lst = tp;
               }
          }
          pq[top[u]].push({srt[u], nw});
          reverse(all(tmp));
          for(auto c: tmp) res += 1ll*get(c.ss-1)*c.ff, upd(c.ss, c.ff);
          u = par[top[u]];
     }
     clear();
     return res;
}
 
int main(){ IOS();
     cin >> n;
     map<int, int> mp; int id = 1;
     rep(i,0,n) cin >> c[i], mp[c[i]] = 0;
     for(auto&c: mp) c.ss = id++;
     rep(i,0,n) c[i] = mp[c[i]];
     vector<pp> edge;
     rep(i,1,n){
          int u, v; cin >> u >> v; u--, v--;
          adj[u].pb(v), adj[v].pb(u);
          edge.pb({u, v});
     }    
     dfs(0, -1), dfsh(0, -1, 0);
     for(auto[u, v]: edge){
          if(h[u] > h[v]) swap(u, v);
          cout << query(u, c[v]) << '\n';          
     }
     return 0-0;
}

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