Submission #1053296

#	Time	Username	Problem	Language	Result	Execution time	Memory
1053296		elazarkoren	Petrol stations (CEOI24_stations)	C++17	55 / 100	845 ms	2097156 KiB

Main

#include <bits/stdc++.h>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define chkmin(a, b) a = min(a, b)
#define chkmax(a, b) a = max(a, b)
#define int ll
using namespace std;
typedef long long ll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef pair<int, int> pii;
typedef vector<pii> vii;
typedef vector<bool> vb;

const int MAX_N = 1e5 + 5;

ll ans[MAX_N];
vii tree[MAX_N];
ll sz[MAX_N];
int n, k;

int DfsSz(int node, int parent) {
    sz[node] = 1;
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent) {
            sz[node] += DfsSz(neighbor, node);
        }
    }
    return sz[node];
}

vvi dp[MAX_N];

vi DpUp(int node, int parent, int uw = 0) {
    for (int i = 0; i < tree[node].size(); i++) {
        if (tree[node][i].x == parent) {
            swap(tree[node][i], tree[node].back());
            break;
        }
    }
    dp[node].resize(tree[node].size());
    vi dp2(k + 1);
    dp2[k - uw] = 1;
    for (int i = 0; i + (parent != -1) < tree[node].size(); i++) {
        auto [neighbor, w] = tree[node][i];
        vi &v = dp[node][i] = DpUp(neighbor, node, w);
        for (int j = 0; j < uw; j++) {
            ans[node] += v[j] * (n - sz[node]);
            dp2[k - uw] += v[j];
        }
        for (int j = uw; j <= k; j++) {
            dp2[j - uw] += v[j];
        }
    }
    return dp2;
}

void DpDown(int node, int parent, vi &dp2) {
    int len = tree[node].size();
    if (parent != -1) dp[node][len - 1] = dp2;
    fill(all(dp2), 0);
    for (int i = 0; i < len; i++) {
        for (int j = 0; j <= k; j++) {
            dp2[j] += dp[node][i][j];
        }
    }
    dp2[k]++;
    for (int i = 0; i + (parent != -1) < len; i++) {
        vi dp3(k + 1);
        auto [neighbor, w] = tree[node][i];
        for (int j = 0; j <= k; j++) {
            dp3[j] = dp2[j] - dp[node][i][j];
        }
        vi nx(k + 1);
        for (int j = 0; j < w; j++) {
            ans[node] += dp3[j] * sz[neighbor];
            nx[k - w] += dp3[j];
        }
        for (int j = w; j <= k; j++) {
            nx[j - w] += dp3[j];
        }
        DpDown(neighbor, node, nx);
//        for (int a = 0; a < len; a++) {
//            if (i == a) continue;
//            vi &v = dp[node][a];
//            for (int j = 0; j < w; j++) {
//                ans[node] += v[j] * sz[neighbor];
//                dp3[k - w] += v[j];
//            }
//            for (int j = w; j <= k; j++) {
//                dp3[j - w] += v[j];
//            }
//        }
//        dp3[k - w]++;
//        DpDown(neighbor, node, dp3);
    }
}

//void Solve(int node, int parent, int d) {
//    for (auto [neighbor, w] : tree[node]) {
//        if (neighbor == parent) continue;
//        if (d < w) {
//            ans[node] += sz[neighbor];
//            Solve(neighbor, node, k - w);
//        } else Solve(neighbor, node, d - w);
//    }
//}

int32_t main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++) {
        int u, v, l;
        cin >> u >> v >> l;
        tree[u].push_back({v, l});
        tree[v].push_back({u, l});
    }
    DfsSz(1, -1);
    DpUp(1, -1);
    vi v(k + 1);
    DpDown(1, -1, v);

//    for (int start = 0; start < n; start++) {
//        if (tree[start].size() != 1) continue;
//        int node = start;
//        vb visited(n);
//        vi cnt(n);
//        queue<pii> q;
//        ll s1 = 0;
//        for (int t = 0; t < n - 1; t++) {
//            visited[node] = 1;
//            int nx = -1, nw;
//            for (auto [neighbor, w] : tree[node]) {
//                if (!visited[neighbor]) {
//                    nx = neighbor;
//                    nw = w;
//                }
//            }
//            q.push({node, s1});
//            s1 += nw;
//            while (!q.empty() && s1 - q.front().y > k) {
//                cnt[node] += cnt[q.front().x];
//                q.pop();
//            }
//            ans[node] += cnt[node] * ll(n - t - 1);
//            cnt[node]++;
//            node = nx;
//        }
//        cout << "";
//    }
//    for (int i = 0; i < n; i++) {
//        DfsSz(i, -1);
//        Solve(i, -1, k);
//    }
    for (int i = 0; i < n; i++) cout << ans[i] << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'vi DpUp(ll, ll, ll)':
Main.cpp:36:23: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   36 |     for (int i = 0; i < tree[node].size(); i++) {
      |                     ~~^~~~~~~~~~~~~~~~~~~
Main.cpp:45:40: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   45 |     for (int i = 0; i + (parent != -1) < tree[node].size(); i++) {
      |                     ~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 18 / 18
• Subtask 2: 8 / 8
• Subtask 3: 0 / 10
• Subtask 4: 12 / 12
• Subtask 5: 17 / 17
• Subtask 6: 0 / 35