Submission #511200

#	Time	Username	Problem	Language	Result	Execution time	Memory
511200		Victor	Janjetina (COCI21_janjetina)	C++17	110 / 110	941 ms	26856 KiB

Main

// #pragma GCC target ("avx,avx2,fma")
// #pragma GCC optimize ("Ofast,inline") // O1 - O2 - O3 - Os - Ofast
// #pragma GCC optimize ("unroll-loops")
#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b) for (int i = (a); i < (b); ++i)
#define per(i, a, b) for (int i = (b - 1); i >= (a); --i)
#define trav(a, x) for (auto &a : x)

#define all(x) x.begin(), x.end()
#define sz(x) x.size()
#define pb push_back
#define debug(x) cout << #x << " = " << x << endl

#define umap unordered_map
#define uset unordered_set

typedef pair<int, int> ii;
typedef pair<int, ii> iii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<vi> vvi;

typedef long long ll;
typedef pair<ll, ll> pll;
typedef vector<ll> vll;
typedef vector<pll> vpll;

const int INF = 1'000'000'007;

struct Node {
    Node *l, *r;
    int lo, hi, val, mset = INF, madd = 0;
    Node(int L, int R) : lo(L), hi(R) {
        val = 0;
        if (hi - lo != 1) {
            int mid = (lo + hi) / 2;
            l = new Node(lo, mid);
            r = new Node(mid, hi);
        }
    }

    int query(int L, int R) {
        if (hi <= L || R <= lo) return 0;
        if (L <= lo && hi <= R) return val;
        push();
        return l->query(L, R) + r->query(L, R);
    }

    void add(int L, int R, int x) {
        if (hi <= L || R <= lo) return;
        if (L <= lo && hi <= R) {
            if (mset != INF)
                mset += x;
            else
                madd += x;

            val += x * (hi - lo);
            return;
        }
        push();
        l->add(L, R, x);
        r->add(L, R, x);
        val = l->val + r->val;
    }

    void set(int L, int R, int x) {
        if (hi <= L || R <= lo) return;
        if (L <= lo && hi <= R) {
            mset = x;
            val = x * (hi - lo);
            madd = 0;
            return;
        }
        push();
        l->set(L, R, x);
        r->set(L, R, x);
        val = l->val + r->val;
    }

    void push() {
        if (mset != INF) {
            l->set(lo, hi, mset);
            r->set(lo, hi, mset);
            mset = INF;

        } else if (madd) {
            l->add(lo, hi, madd);
            r->add(lo, hi, madd);
            madd = 0;
        }
    }
};

bitset<100005> removed;
Node *segtree;
int n, k, subsize[100005], centroid, cnt = 0;
vii graph[100005];

void dfs(int u, int p) {
    subsize[u] = 1;
    int mx = 0;
    trav(edge, graph[u]) {
        int v, w;
        tie(v, w) = edge;
        if (v != p && !removed[v]) {
            dfs(v, u);
            subsize[u] += subsize[v];
            mx = max(mx, subsize[v]);
        }
    }

    if (subsize[u] >= (cnt + 1) / 2 && mx <= cnt / 2) centroid = u;
}

ll pairs(vii &paths) {
    sort(all(paths));
    segtree->set(0, n, 0);

    ll ans = 0;
    trav(path, paths) {
        int mx, dist;
        tie(mx, dist) = path;
        // cout << "mx = " << mx << " dist = " << dist << " pairs = " << segtree->query(0, mx - dist - k) << endl << endl;
        ans += segtree->query(0, mx - dist - k);
        segtree->add(dist, dist + 1, 1);
    }

    return ans;
}

ll ans = 0;
vii paths, subpaths;
void dfs2(int u, int p, int mx_road, int dist) {
    paths.emplace_back(mx_road, dist);
    if (u != centroid) subpaths.emplace_back(mx_road, dist);

    trav(edge, graph[u]) {
        int v, w;
        tie(v, w) = edge;
        if (v != p && !removed[v]) {
            dfs2(v, u, max(mx_road, w), dist + 1);
            if (u == centroid) {
                ans -= pairs(subpaths);
                subpaths.clear();
            }
        }
    }
}

void predfs(int u,int p) {
    ++cnt;
    trav(edge,graph[u]){
        int v=edge.first;
        if(v!=p&&!removed[v])predfs(v,u);
    }
}

void centroid_decomp(int u) {
    cnt=0;
    predfs(u,u);
    dfs(u, u);
    if (subsize[u] == 1) return;
    assert(!removed[centroid]);

    dfs2(centroid, centroid, 0, 0);
    ans += pairs(paths);
    paths.clear();

    removed[centroid] = 1;
    trav(edge, graph[centroid]) {
        int v, w;
        tie(v, w) = edge;
        if (!removed[v]) centroid_decomp(v);
    }
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    cin.exceptions(cin.failbit);

    cin >> n >> k;
    segtree = new Node(0, n);
    --k;
    rep(i, 1, n) {
        int u, v, w;
        cin >> u >> v >> w;
        graph[--u].emplace_back(--v, w);
        graph[v].emplace_back(u, w);
    }

    centroid_decomp(0);
    cout << ans * 2 << endl;
    return 0;
}

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