4948 2023. 04. 07 18:52:52 zsombor Testvérvárosok cpp17 Accepted 100/100 180ms 17196 KiB
#include <iostream>
#include <vector>
using namespace std;
using ll = long long;

struct edge {
    int to, w;
    edge(int TO, int W) {
        to = TO;
        w = W;
    }
};

int n, k, c, sz;
ll wsum, ans = 0;
vector <bool> volt(5e4 + 1, false);
vector <vector <edge>> g(5e4 + 1);
vector <int> cnt(1e3, 0);

int get_c(int x) {
    if (volt[x]) return 0;
    volt[x] = true;
    int cur, sum = 1;
    bool cen = true;
    for (edge e : g[x]) {
        cur = get_c(e.to);
        sum += cur;
        if (cur > sz / 2) cen = false;
    }
    if (sz - sum > sz / 2) cen = false;
    if (cen) c = x;
    volt[x] = false;
    return sum;
}

void dfs1(int x) {
    if (volt[x]) return;
    volt[x] = true;
    ans += cnt[(k - wsum % k) % k];
    for (edge e : g[x]) {
        wsum += e.w;
        dfs1(e.to);
        wsum -= e.w;
    }
    volt[x] = false;
}

void dfs2(int x) {
    if (volt[x]) return;
    volt[x] = true;
    cnt[wsum % k]++;
    for (edge e : g[x]) {
        wsum += e.w;
        dfs2(e.to);
        wsum -= e.w;
    }
    volt[x] = false;
}

void solve(int x) {
    if (volt[x]) return;
    sz = get_c(x);
    get_c(x);
    volt[c] = true;
    fill(cnt.begin(), cnt.end(), 0);
    cnt[0] = 1;
    for (edge e : g[c]) {
        wsum = e.w;
        dfs1(e.to);
        dfs2(e.to);
    }
    int C = c;
    for (edge e : g[C]) solve(e.to);
}

int main()
{
    cin >> n >> k;
    int a, b, W;
    for (int i = 0; i < n - 1; i++) {
        cin >> a >> b >> W;
        g[a].push_back(edge(b, W));
        g[b].push_back(edge(a, W));
    }
    solve(1);
    cout << ans;
}
Subtask Sum Test Verdict Time Memory
subtask1 0/0
1 Accepted 3ms 3920 KiB
2 Accepted 3ms 4260 KiB
subtask2 15/15
3 Accepted 3ms 4392 KiB
4 Accepted 3ms 4456 KiB
5 Accepted 4ms 4348 KiB
6 Accepted 4ms 4600 KiB
7 Accepted 4ms 4556 KiB
8 Accepted 4ms 4864 KiB
9 Accepted 4ms 4812 KiB
subtask3 15/15
10 Accepted 39ms 5656 KiB
11 Accepted 10ms 5652 KiB
12 Accepted 171ms 8420 KiB
13 Accepted 155ms 8420 KiB
14 Accepted 180ms 8900 KiB
subtask4 20/20
15 Accepted 4ms 5204 KiB
16 Accepted 4ms 5368 KiB
17 Accepted 18ms 6436 KiB
18 Accepted 68ms 10328 KiB
19 Accepted 98ms 12676 KiB
20 Accepted 137ms 15616 KiB
21 Accepted 128ms 15212 KiB
22 Accepted 156ms 17196 KiB
subtask5 50/50
23 Accepted 158ms 9656 KiB
24 Accepted 82ms 8152 KiB
25 Accepted 114ms 8580 KiB
26 Accepted 25ms 6960 KiB
27 Accepted 82ms 8256 KiB
28 Accepted 171ms 9904 KiB
29 Accepted 159ms 9844 KiB
30 Accepted 159ms 9848 KiB
31 Accepted 180ms 10272 KiB