Submission #486233

#	Time	Username	Problem	Language	Result	Execution time	Memory
486233		jalsol	Janjetina (COCI21_janjetina)	C++17	110 / 110	347 ms	14284 KiB

Main

#ifdef LOCAL
#include "local.h"
#else
#include <bits/stdc++.h>
#define debug(...)
#define DB(...)
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#endif

using namespace std;

static const bool __initialization = []() {
    cin.tie(nullptr)->sync_with_stdio(false);
#define TASK "JANJETINA"
    if (fopen(TASK".inp", "r")) {
        (void)(!freopen(TASK".inp", "r", stdin));
        (void)(!freopen(TASK".out", "w", stdout)); }
    return true;
}();

using ll = long long;
using ld = long double;
#define fi first
#define se second

#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()

#define For(i, l, r) for (int i = (l); i <= (r); ++i)
#define Ford(i, r, l) for (int i = (r); i >= (l); --i)
#define Rep(i, n) For (i, 0, (n) - 1)
#define Repd(i, n) Ford (i, (n) - 1, 0)
#define Fe(...) for (auto __VA_ARGS__)

template <class C> inline int isz(const C& c) { return static_cast<int>(c.size()); }
template <class T> inline bool chmin(T& a, const T& b) { return (a > b) ? a = b, true : false; }
template <class T> inline bool chmax(T& a, const T& b) { return (a < b) ? a = b, true : false; }

constexpr ld eps = 1e-9;
constexpr int inf = 1e9;
constexpr ll linf = 1e18;

// =============================================================================

constexpr int maxn = 1e5 + 5;

struct Fen {
    int fen[maxn];

    Fen() {
        memset(fen, 0, sizeof(fen));
    }

    void update(int i, int x) {
        for (; i < maxn; i += (i & -i)) {
            fen[i] += x;
        }
    }

    int get(int i) {
        int ret = 0;
        for (; i > 0; i -= (i & -i)) {
            ret += fen[i];
        }
        return ret;
    }
};

struct Link { int v, w; };

int n, k;
vector<Link> g[maxn];

int sz[maxn];
bool rem[maxn];

int getSz(int u, int p = -1) {
    sz[u] = 1;

    Fe (&[v, _] : g[u]) {
        if (!rem[v] && v != p) {
            sz[u] += getSz(v, u);
        }
    }

    return sz[u];
}

int getCen(int target, int u, int p = -1) {
    Fe (&[v, _] : g[u]) {
        if (rem[v] || v == p) continue;

        if (sz[v] * 2 > target) {
            return getCen(target, v, u);
        }
    }

    return u;
}

vector<pair<int, int>> other;
Fen fen;
ll ans = 0;

void findSub(int u, int p, int maxw, int dep) {
    other.emplace_back(maxw, dep);

    Fe (&[v, w] : g[u]) {
        if (!rem[v] && v != p) {
            findSub(v, u, max(maxw, w), dep + 1);
        }
    }
}

ll calc() {
    ll ret = 0;
    sort(all(other));

    // max(w1, w2) - (l1 + l2) >= k

    // max(w1, w2) = w2
    // <=> w2 - l1 - l2 >= k
    // <=> w2 - l2 >= k + l1

    Fe (&[maxw, dep] : other) {
        ret += fen.get(maxw - dep);
        fen.update(dep + k, 1);
    }

    Fe (&[maxw, dep] : other) {
        fen.update(dep + k, -1);
    }

    return ret;
}

void solve(int u) {
    u = getCen(getSz(u), u);
    rem[u] = true;

    other.clear();
    findSub(u, -1, 0, 0);
    ans += calc();

    // only consider paths on different subtrees
    // => eliminate paths on the same subtree
    Fe (&[v, w] : g[u]) {
        if (!rem[v]) {
            other.clear();
            findSub(v, u, w, 1);
            ans -= calc();
        }
    }

    Fe (&[v, _] : g[u]) {
        if (!rem[v]) {
            solve(v);
        }
    }
}

signed main() {
    cin >> n >> k;
    Rep (_, n - 1) {
        int u, v, w; cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    solve(1);

    cout << ans * 2 << '\n';
}

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