Submission #1083974

#	Time	Username	Problem	Language	Result	Execution time	Memory
1083974		tfgs	Janjetina (COCI21_janjetina)	C++17	0 / 110	1 ms	348 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <class T> using Tree = tree<T, null_type, less_equal<T>, 
    rb_tree_tag, tree_order_statistics_node_update>;
#define int int64_t

#ifdef LOCAL
#include "algo/debug.h"
#endif

#define f first
#define s second
template<class T> using V = vector<T>; 
using vi = V<int>;

#define all(x) begin(x), end(x)
#define rall(x) rbegin(x), rend(x) 
#define pb push_back
#define lb lower_bound
#define ub upper_bound
template<class T> int lwb(V<T>& a, const T& b) { return lb(all(a),b)-begin(a); }
template<class T> int upb(V<T>& a, const T& b) { return ub(all(a),b)-begin(a); }
template<class T> bool ckmin(T& a, const T& b) { return a > b ? a=b, true : false; }
template<class T> bool ckmax(T& a, const T& b) { return a < b ? a=b, true : false; }


int n;
vi sub;
V<bool> del;
int ans;
V<V<pair<int, array<int, 2>>>> g;
V<pair<array<int, 2>, int>> paths;

void find_sub(int u, int p) {
    sub[u] = 1;
    for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
        find_sub(v, u);
        sub[u] += sub[v];
    }
}

int find_ctd(int u, int p, int sub_siz) {
    for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
        if (2*sub[v] >= sub_siz) {
            return find_ctd(v, u, sub_siz);
        }
    }
    return u;
}

void find_paths(int u, int p, int d, array<int, 2> mx) {
    for (auto [v, wi] : g[u]) if (v != p && !del[v]) {
        ckmax(mx, wi);
        paths.pb({mx, d+1});
        find_paths(v, u, d+1, mx);
        if (p == -1) {
            mx = {0, -1};
        }
    }
}

int count_pairs_in_sub(int u, int p) {
    int res = 0;
    find_paths(u, p, 0, {0, -1});
    sort(all(paths));
    Tree<int> s;
    for (auto [wi, d] : paths) {
        int mx = wi[0];
        if (mx >= d) res++;
        res += s.order_of_key(mx-d+1);
        s.insert(d);
    }
    paths.clear();
    return res;
}

void process_ctd(int ctd, int p) {
    // ps("processing centroid", ctd+1);
    // ps("+", count_pairs_in_sub(ctd, p));
    ans += count_pairs_in_sub(ctd, p);
    for (auto [v, wi] : g[ctd]) if (v != p && !del[v]) {
        // ps("-", count_pairs_in_sub(v, ctd));
        ans -= count_pairs_in_sub(v, ctd);
    }
}

void cd(int u) {
    find_sub(u, -1);
    int ctd = find_ctd(u, -1, sub[u]);
    process_ctd(ctd, -1);
    del[ctd] = true;
    for (auto [v, wi] : g[ctd]) if (!del[v]) {
        cd(v);
    }
}

void solve() {
    int k; cin >> n >> k; 
    g.resize(n);
    del.resize(n);
    sub.resize(n);

    for (int i = 0; i < n-1; i++) {
        int u, v, w; cin >> u >> v >> w; u--; v--; w -= k;
        g[u].pb({v, {w, i}}); g[v].pb({u, {w, i}});
    }
    cd(0);
    cout << 2*ans << '\n';
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    solve();
    return 0;
}

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