Submission #785881

#	Time	Username	Problem	Language	Result	Execution time	Memory
785881		RushB	Janjetina (COCI21_janjetina)	C++17	110 / 110	356 ms	24764 KiB

Main

#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:31: warning: unused variable 'w' [-Wunused-variable]
   86 |                 int u = P[0], w = P[1];
      |                               ^

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60