Submission #387202

#	Time	Username	Problem	Language	Result	Execution time	Memory
387202		phathnv	Janjetina (COCI21_janjetina)	C++11	15 / 110	8 ms	5356 KiB

Main

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int N = 1e5 + 7;

struct Edge{
    int v, c;
    Edge(int _v, int _c){
        v = _v;
        c = _c;
    }
};

int n, k, sz[N];
vector<Edge> adj[N];
ll answer = 0;
bool done[N];

void DfsForCentroid(int u, int p){
    sz[u] = 1;
    for(Edge e : adj[u]){
        int v = e.v;
        if (v == p || done[v])
            continue;
        DfsForCentroid(v, u);
        sz[u] += sz[v];
    }
}

int FindCentroid(int u, int p, int treeSize){
    for(Edge e : adj[u]){
        int v = e.v;
        if (v == p || done[v])
            continue;
        if (sz[v] * 2 > treeSize)
            return FindCentroid(v, u, treeSize);
    }
    return u;
}

void Dfs(int u, int p, int len, int maxW, vector<pair<int, int>> &a){
    if (done[u])
        return;

    a.push_back({maxW, len});
    for(Edge e : adj[u]){
        int v = e.v;
        int c = e.c;
        if (v == p)
            continue;
        Dfs(v, u, len + 1, max(maxW, c), a);
    }
}

ll Calc(vector<pair<int, int>> &a){
    ll res = 0;
    for(int i = 0; i < (int) a.size(); i++)
        for(int j = i + 1; j < (int) a.size(); j++){
            int l = a[i].second + a[j].second;
            int w = max(a[i].first, a[j].first);
            if (w - l >= k)
                res++;
        }
    return res;
}

void Solve(int u){
    if (done[u])
        return;

    DfsForCentroid(u, -1);
    u = FindCentroid(u, -1, sz[u]);
    vector<pair<int, int>> a;
    Dfs(u, -1, 0, 0, a);
    answer += Calc(a);
    for(Edge e : adj[u]){
        a.clear();
        int v = e.v;
        int c = e.c;
        Dfs(v, u, 1, c, a);
        answer -= Calc(a);
    }

    done[u] = 1;
    for(Edge e : adj[u])
        Solve(e.v);
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> k;
    assert(n <= 1000);
    for(int i = 1; i < n; i++){
        int u, v, c;
        cin >> u >> v >> c;
        adj[u].push_back(Edge(v, c));
        adj[v].push_back(Edge(u, c));
    }

    Solve(1);
    cout << answer * 2;

    return 0;
}

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