Submission #1159842

#TimeUsernameProblemLanguageResultExecution timeMemory
1159842itslqJanjetina (COCI21_janjetina)C++20
15 / 110
1600 ms142148 KiB
#include <bits/stdc++.h>
using namespace std;

#define int long long

const int MAX = 4e5, LOG = 18;
vector<int> dist(MAX);
vector<bool> vis(MAX);
vector<vector<pair<int, int>>> adjList(MAX);
vector<vector<pair<int, int>>> pa(MAX, vector<pair<int, int>>(LOG, make_pair(-1, -1)));

pair<int, int> kpa(int n, int k) {
    int maxw = 0;
    for (int i = 0; i < LOG; i++) {
        if ((1 << i) & k) {
            maxw = max(maxw, pa[n][i].second);
            n = pa[n][i].first;
        }
    }
    return make_pair(n, maxw);
}

int ans(int a, int b) {
    int maxw, S = 0;

    if (dist[a] < dist[b]) swap(a, b);
    pair<int, int> res = kpa(a, S = dist[a] - dist[b]);
    a = res.first;
    maxw = res.second;

    if (a == b) return maxw - S;

    for (int i = LOG - 1; i >= 0; i--) {
        if (pa[a][i].first != pa[b][i].first && pa[a][i].first != -1 && pa[b][i].first != -1) {
            maxw = max(maxw, max(pa[a][i].second, pa[b][i].second));
            a = pa[a][i].first;
            b = pa[b][i].first;
            S += (1 << (i + 1));
        }
    }

    return max(maxw, max(pa[a][0].second, pa[b][0].second)) - S - 2;
}

signed main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	
    int N, K, Q, A, B, W, T = 0;
    pair<pair<int, int>, int> curr;
    vector<pair<pair<int, int>, int>> edges;
    cin >> N >> K;
    for (int i = 0; i < N - 1; i++) {
        cin >> A >> B >> W;
        edges.emplace_back(make_pair(--A, --B), W);
        adjList[A].push_back(make_pair(W, B));
        adjList[B].push_back(make_pair(W, A));
    }

    vis[0] = 1;
    priority_queue<pair<pair<int, int>, int>, vector<pair<pair<int, int>, int>>, greater<pair<pair<int, int>, int>>> pq;
    for (pair<int, int> adj: adjList[0]) pq.push(make_pair(adj, 0));

    while (!pq.empty()) {
        curr = pq.top();
        pq.pop();

        if (vis[curr.first.second]) continue;
        vis[curr.first.second] = 1;
        pa[curr.first.second][0] = make_pair(curr.second, curr.first.first);
        dist[curr.first.second] = dist[curr.second] + 1;
        T += curr.first.first;

        for (pair<int, int> adj: adjList[curr.first.second]) {
            if (!vis[adj.second]) {
                pq.push(make_pair(adj, curr.first.second));
            }
        }
    }

    for (int j = 1; j < LOG; j++) {
	    for (int i = 1; i < N; i++) {
			if (pa[i][j - 1].first == -1) continue;
            pa[i][j].first = pa[pa[i][j - 1].first][j - 1].first;
            pa[i][j].second = max(pa[i][j - 1].second, pa[pa[i][j - 1].first][j - 1].second);
        }
    }
    
    A = 0;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < i; j++) {
            // cout << i << "-" << j << "-" << ans(i, j) << endl;
            if (ans(i, j) >= K) A++;
        }
    }
    
    cout << (A << 1);
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...