Submission #922124

#TimeUsernameProblemLanguageResultExecution timeMemory
922124406Janjetina (COCI21_janjetina)C++17
110 / 110
353 ms25012 KiB
    #include <bits/stdc++.h>
    using namespace std;
    #define FOR(i, a, b) for (int i = (a); i < (b); i++) 
    #define int long long
    using ar = array<int, 2>;
    const int64_t INF = 1ll << 60;
    const int N = 1e5 + 5;
    vector<bool> del(N);
    int sub[N], ptr, M[N], h[N], k, n, bit[N], ans;
    vector<ar> adj[N], VV, V[N];
     
    void add(int r, int d) {
            r++;
            for (; r < N; r += r & -r) bit[r] += d;
    }
    int get(int r) {
            r++;
            int re = 0;
            for (; r > 0; r &= (r - 1)) re += bit[r];
            return re;
    }
     
    void find_sub(int v, int p) {
            sub[v] = 1;
            for (auto P: adj[v]) {
                    int u = P[0];
                    if (del[u] || u == p) continue;
                    find_sub(u, v);
                    sub[v] += sub[u];
            }
    }
    int find_center(int v, int p, int kol) {
            for (auto P: adj[v]) {
                    int u = P[0];
                    if (del[u] || u == p || sub[u] <= (kol / 2)) continue;
                    return find_center(u, v, kol);
            }
            return v;
    }
    void dfs(int v, int p) {
            for (auto P: adj[v]) {
                    int u = P[0], w = P[1];
                    if (u == p || del[u]) continue;
                    M[u] = max(M[v], w);
                    h[u] = h[v] + 1;
                    dfs(u, v);
            }
            V[ptr].push_back({M[v], h[v]});
            VV.push_back({M[v], h[v]});
    }
     
    int solve(vector<ar> &a) {
            sort(a.begin(), a.end());
            int ans = 0;
            for (auto P: a) {
                    int m = P[0], h = P[1];
                    ans += get(m - h - k);
                    add(h, +1);
            }
            for (auto P: a)
                    add(P[1], -1);
            return ans;
    }
     
    void rec(int x) {
            find_sub(x, x);
            int c = find_center(x, x, sub[x]);
            del[c] = true;
            for (auto P: adj[c]) {
                    int u = P[0], w = P[1];
                    if (del[u]) continue;
                    h[u] = 1, M[u] = w;
                    dfs(u, c);
                    ptr++;
            }
            FOR(i, 0, ptr) {
                    ans -= solve(V[i]);
                    V[i].clear();
            }
            VV.push_back({-INF, 0});
            ans += solve(VV);
            VV.clear();
            ptr = 0;
     
            for (auto P: adj[c]) {
                    int u = P[0], w = P[1];
                    if (del[u]) continue;
                    rec(u);
            }
    }
    signed main() {
            ios::sync_with_stdio(0), cin.tie(0);
            cin >> n >> k;
            FOR(i, 1, n) {
                    int u, v, w; 
                    cin >> u >> v >> w;
                    u--, v--;
                    adj[v].push_back({u, w});
                    adj[u].push_back({v, w});
            }
            rec(0);
            cout << 2ll * ans << '\n';
    }
    //20:55:05

Compilation message (stderr)

Main.cpp: In function 'void rec(long long int)':
Main.cpp:86:35: warning: unused variable 'w' [-Wunused-variable]
   86 |                     int u = P[0], w = P[1];
      |                                   ^
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...