Submission #515599

#	Time	Username	Problem	Language	Result	Execution time	Memory
515599		KoD	Janjetina (COCI21_janjetina)	C++17	110 / 110	401 ms	14888 KiB

Main

#include <bits/stdc++.h>

using std::vector;
using std::array;
using std::pair;
using std::tuple;

template <class F> struct RecLambda : private F {
    explicit RecLambda(F&& f) : F(std::forward<F>(f)) {}
    template <class... Args> decltype(auto) operator()(Args&&... args) const {
        return F::operator()(*this, std::forward<Args>(args)...);
    }
};

using ll = long long;

struct Fenwick {
    int size;
    vector<int> data;
    explicit Fenwick(const int n) : size(n), data(n + 1) {}
    void add(int i, const int x) {
        i += 1;
        while (i <= size) {
            data[i] += x;
            i += i & -i;
        }
    }
    int pref(int i) const {
        i += 1;
        int ret = 0;
        while (i > 0) {
            ret += data[i];
            i -= i & -i;
        }
        return ret;
    }
};

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int N, K;
    std::cin >> N >> K;
    vector<vector<pair<int, int>>> graph(N);    
    for (int i = 1; i < N; ++i) {
        int x, y, w;
        std::cin >> x >> y >> w;
        x -= 1, y -= 1;
        graph[x].emplace_back(y, w);
        graph[y].emplace_back(x, w);
    }
    Fenwick fen(N);
    const auto count = [&](vector<pair<int, int>>& yx) {
        std::sort(yx.begin(), yx.end());
        ll ret = 0;
        for (const auto& [y, x] : yx) {
            ret += fen.pref(std::min(N - 1, y - x - K));
            fen.add(x, 1);
        }
        for (const auto& [y, x] : yx) {
            fen.add(x, -1);
        }
        return ret;
    };
    vector<char> dead(N);
    vector<int> subtree(N);
    ll ans = 0;
    RecLambda([&](auto&& dfs, const int root) -> void {
        RecLambda([&](auto&& dfs, const int u, const int p) -> void {
            subtree[u] = 1;
            for (const auto& [v, w] : graph[u]) {
                if (!dead[v] and v != p) {
                    dfs(v, u);
                    subtree[u] += subtree[v];
                }
            }
        })(root, -1);
        const int cent = RecLambda([&](auto&& dfs, const int u, const int p) -> int {
            for (const auto& [v, w] : graph[u]) {
                if (!dead[v] and v != p and subtree[v] * 2 > subtree[root]) {
                    return dfs(v, u);
                }
            }
            return u;
        })(root, -1);
        vector<pair<int, int>> list;
        vector<pair<int, int>> all;
        const auto enumerate = RecLambda([&](auto&& dfs, const int u, const int p, const int x, const int y) -> void {
            if (y - x >= K) {
                ans += 1;
            }
            list.emplace_back(y, x);
            all.emplace_back(y, x);
            for (const auto& [v, w] : graph[u]) {
                if (!dead[v] and v != p) {
                    dfs(v, u, x + 1, std::max(y, w));
                }
            }
        });
        for (const auto& [u, w] : graph[cent]) {
            if (!dead[u]) {
                list.clear();
                enumerate(u, cent, 1, w);
                ans -= count(list);
            }
        }
        ans += count(all);
        dead[cent] = true;
        for (const auto& [u, w] : graph[cent]) {
            if (!dead[u]) {
                dfs(u);
            }
        }
    })(0);
    std::cout << ans * 2 << '\n';
    return 0;
}

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