Submission #1220342

#	Time	Username	Problem	Language	Result	Execution time	Memory
1220342		LucaIlie	Petrol stations (CEOI24_stations)	C++20	18 / 100	3594 ms	18612 KiB

Main

#include <bits/stdc++.h>

using namespace std;

struct edge {
    int u, v, c;

    int other(int w) {
        return u ^ v ^ w;
    }
};

struct SEG_TREE {
    int ls, rs;
    int root;
    vector<int> segTree, leftSon, rightSon;
    vector<int> brut;

    void init(int l, int r) {
        ls = l;
        rs = r;
        segTree.resize(1);
        leftSon.resize(1);
        rightSon.resize(1);
        root = 0;
        segTree[root] = 0;
        leftSon[root] = rightSon[root] = -1;
        brut.clear();
        brut.resize(r + 1);
    }

    int createNode() {
        int v = segTree.size();
        segTree.push_back(0);
        leftSon.push_back(-1);
        rightSon.push_back(-1);
        return v;
    }

    void update(int v, int l, int r, int p, int x) {
        if (l > p || r < p)
            return;

        if (l == r) {
            segTree[v] += x;
            return;
        }

        if (leftSon[v] == -1)
            leftSon[v] = createNode();
        if (rightSon[v] == -1)
            rightSon[v] = createNode();

        int mid = (l + r) / 2;
        update(leftSon[v], l, mid, p, x);
        update(rightSon[v], mid + 1, r, p, x);
        segTree[v] = segTree[leftSon[v]] + segTree[rightSon[v]];
    }

    int query(int v, int l, int r, int lq, int rq) {
        if (v == -1)
            return 0;

        if (l > rq || r < lq)
            return 0;

        if (lq <= l && r <= rq)
            return segTree[v];

        int mid = (l + r) / 2;
        return query(leftSon[v], l, mid, lq, rq) + query(rightSon[v], mid + 1, r, lq, rq);
    }

    void modify(int p, int x) {
        // brut[p] += x;
        update(root, ls, rs, p, x);
    }

    int countLower(int p) {
        // int cnt = 0;
        // for (int i = 0; i < p; i++)
        //     cnt += brut[i];
        // return cnt;
        p--;
        return query(root, ls, rs, 0, p);
    }
};

const int MAX_N = 7e4;
int k;
long long ans[MAX_N];
bool isCentroid[MAX_N];
int sizee[MAX_N], depth[MAX_N], starts[MAX_N];
int remUp[MAX_N];
edge edges[MAX_N];
vector<int> adj[MAX_N];
SEG_TREE frecv;

void calcSizes(int u, int p) {
    sizee[u] = 1;
    for (int e: adj[u]) {
        int v = edges[e].other(u);
        if (v == p || isCentroid[v])
            continue;
        calcSizes(v, u);
        sizee[u] += sizee[v];
    }
}

int totsz, centroid;
void findCentroid(int u, int p) {
    int maxsz = totsz - sizee[u];
    for (int e: adj[u]) {
        int v = edges[e].other(u);
        if (v == p || isCentroid[v])
            continue;
        findCentroid(v, u);
        maxsz = max(maxsz, sizee[v]);
    }
    if (maxsz <= totsz / 2)
        centroid = u;
}

vector<int> vert, st, vertBySubtree[MAX_N];
void calcUp(int u, int p, int w) {
    if (p != -1 && w == -1)
        w = u;

    vert.push_back(u);
    st.push_back(u);
    if (p != -1)
        vertBySubtree[w].push_back(u);

    int l = -1, r = st.size() - 1;
    while (r - l > 1) {
        int mid = (l + r) / 2;
        if (depth[u] - depth[st[mid]] <= k)
            r = mid;
        else
            l = mid;
    }
    remUp[u] = (r == 0 ? k - depth[u] : remUp[st[r]]);

    sizee[u] = 1;
    for (int e: adj[u]) {
        int v = edges[e].other(u);
        if (v == p || isCentroid[v])
            continue; 
        depth[v] = depth[u] + edges[e].c;
        calcUp(v, u, w);
        sizee[u] += sizee[v];
    }

    if (r != 0)
        starts[st[r]] += starts[u] + 1;    


    // printf("%d %lld %d %d\n", u, ans[u], sizee[u], sizee[w]);

    st.pop_back();
}

void calcDown(int u, int p) {
    int st = frecv.countLower(depth[u]) - frecv.countLower(depth[p]);
    // printf("ciic %d %d %d %d\n", u, p, st, depth[u]);
    ans[p] += (long long)st * sizee[u];
    frecv.modify(k + depth[p], st);

    for (int e: adj[u]) {
        int v = edges[e].other(u);
        if (v == p)
            continue;
        calcDown(v, u);
    }

    frecv.modify(k + depth[p], -st);
}

void decomp(int r) {
    calcSizes(r, -1);
    totsz = sizee[r];
    findCentroid(r, -1);
    int c = centroid;

    calcUp(c, -1, -1);

    // printf("CEN %d\n", c);
    // for (int v: vert)
    //     printf("%d %d\n", v, remUp[v]);

    for (int e: adj[c]) {
        int w = edges[e].other(c);
        if (isCentroid[w])
            continue;
        for (int v: vertBySubtree[w])
            ans[v] += (long long)(sizee[c] - sizee[w]) * starts[v];
        // printf("%d %d\n", sizee[c], sizee[w]);
    }

    frecv.init(0, (long long)sizee[c] * k);
    for (int v: vert)
        frecv.modify(remUp[v], 1);

    // for (int v: vert)
    //     printf("%d %lld\n", v, ans[v]);

    for (int e: adj[c]) {
        int w = edges[e].other(c);
        if (isCentroid[w])
            continue;

        for (int v: vertBySubtree[w])
            frecv.modify(remUp[v], -1);

        calcDown(w, c);

        for (int v: vertBySubtree[w])
            frecv.modify(remUp[v], 1);

    }

    for (int v: vert) {
        starts[v] = sizee[v] = depth[v] = 0;
        vertBySubtree[v].clear();
    }
    vert.clear();

    isCentroid[c] = true;
    for (int e: adj[c]) {
        int w = edges[e].other(c);
        if (isCentroid[w])
            continue;
        decomp(w);
    }
}

int main() {
    int n;
    
    cin >> n >> k;
    for (int e = 0; e < n - 1; e++) {
        cin >> edges[e].u >> edges[e].v >> edges[e].c;
        adj[edges[e].u].push_back(e);
        adj[edges[e].v].push_back(e);
    }

    frecv.init(0, (long long)n * k);
    decomp(0);

    for (int v = 0; v < n; v++)
        cout << ans[v] << "\n";

    return 0;
}

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