Submission #1177722

#	Time	Username	Problem	Language	Result	Execution time	Memory
1177722		RandomUser	Janjetina (COCI21_janjetina)	C++20	110 / 110	311 ms	21696 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#define int long long
using pii = pair<int, int>;

const int maxn = 2e5 + 5;

struct fenwick {
    int n;
    vector<int> tree;

    fenwick(int _n) : n(_n+10), tree(n+50) {}

    void update(int p, int v) {
        for(p++; p<n; p+=p&-p) tree[p] += v;
    }

    int query(int p) {
        int ans = 0;
        if(p < 0) return ans;
        for(p++; p; p-=p&-p) ans += tree[p];
        return ans;
    }
} fwt(maxn);

long long ans = 0;
int n, del[maxn], sub[maxn], mx[maxn], k, dep[maxn];
vector<int> ord;
vector<pii> G[maxn];

int get_sub(int u, int p) {
    sub[u] = 1;
    for(auto [v, _] : G[u])
        if(v != p && !del[v]) sub[u] += get_sub(v, u);
    return sub[u];
}

int get_cen(int u, int p, int sz) {
    for(auto [v, _] : G[u])
        if(v != p && !del[v] && 2 * sub[v] > sz)
            return get_cen(v, u, sz);
    return u;
}

void dfs(int u, int p) {
    ord.push_back(u);
    for(auto [v, w] : G[u]) {
        if(v == p || del[v]) continue;
        dep[v] = dep[u] + 1;
        mx[v] = max(mx[u], w);
        dfs(v, u);
    }
}

bool cmp(int u, int v) {
    return mx[u] < mx[v];
}

int calc() { 
    int cnt = 0;
    sort(ord.begin(), ord.end(), cmp);

    for(int u : ord) {
        cnt += fwt.query(mx[u]-dep[u]-k);
        fwt.update(dep[u], 1);
    }

    for(int u : ord) 
        fwt.update(dep[u], -1);
    return cnt;
}

void build(int u) {
    int cen = get_cen(u, u, get_sub(u, u));
    del[cen] = 1;

    ord.clear();
    mx[cen] = dep[cen] = 0;
    dfs(cen, cen);

    ans += calc();
    for(auto [v, w] : G[cen]) {
        if(del[v]) continue;
        ord.clear();
        mx[v] = w;
        dep[v] = 1;
        dfs(v, cen);
        ans -= calc();
    }

    for(auto [v, _] : G[cen]) if(!del[v]) build(v);
}

signed main() {
    cin >> n >> k;
    for(int i=1; i<n; i++) {
        int a, b, w; cin >> a >> b >> w;
        G[a].push_back({ b, w });
        G[b].push_back({ a, w });
    }

    build(1);
    cout << 2 * ans << '\n';
}

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