Submission #245250

#	Time	Username	Problem	Language	Result	Execution time	Memory
245250		Tc14	Railway (BOI17_railway)	C++17	100 / 100	249 ms	39984 KiB

railway

//#pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;
#define ve vector

struct SegmentTree {

    int Cap;
    ve<int> Seg;

    void build(int n) {
        Cap = 1 << (int)ceil(log2(n));
        Seg = ve<int>(2 * Cap);
    }

    int left(int i) { return 2 * i; }
    int right(int i) { return 2 * i + 1; }
    int parent(int i) { return i / 2; }

    int query(int x) {

        int curr, r;

        curr = Cap + x;
        r = 0;
        while (curr != 0) {
            r += Seg[curr];
            curr = parent(curr);
        }

        return r;
    }

    void rangeUpdate(int a, int b, int x) {
        update(a, b, 0, Cap - 1, x, 1);
    }

    void update(int a, int b, int i, int j, int x, int curr) {

        int m;

        if (a <= i && j <= b) {
            Seg[curr] += x;
            return;
        }
        if (j < a || b < i) return;

        m = (i + j) / 2;

        update(a, b, i, m, x, left(curr));
        update(a, b, m + 1, j, x, right(curr));
    }
};

int L, C;
ve<int> Depth, Size, Jobs;
ve<pair<int, int>> Comp;
ve<ve<pair<int, int>>> Graph, Intervals;
ve<ve<int>> Ancestor, CompVertecies, CompEdges;
ve<SegmentTree> SegTrees;

void dfs(int u, int p, int d) {

    int v, a;

    Depth[u] = d;

    a = p;
    for (int i = 0; i < L; i++) {
        Ancestor[u][i] = a;
        if (a != -1) a = Ancestor[a][i];
    }

    Size[u] = 1;
    for (pair<int, int> e : Graph[u]) {
        v = e.first;
        if (v != p) {
            dfs(v, u, d + 1);
            Size[u] += Size[v];
        }
    }
}

tuple<int, int, int> lca(int a, int b) {

    int ha, hb;

    ha = 0;
    hb = 0;

    if (Depth[a] > Depth[b]) {
        for (int i = L - 1; i >= 0; i--) {
            if (Depth[a] - (1 << i) >= Depth[b]) {
                a = Ancestor[a][i];
                ha += 1 << i;
            }
        }
    }
    else {
        for (int i = L - 1; i >= 0; i--) {
            if (Depth[b] - (1 << i) >= Depth[a]) {
                b = Ancestor[b][i];
                hb += 1 << i;
            }
        }
    }

    if (a == b) return {a, ha, hb};
    else {
        for (int i = L - 1; i >= 0; i--) {
            if (Ancestor[a][i] != Ancestor[b][i]) {
                a = Ancestor[a][i];
                b = Ancestor[b][i];
                ha += 1 << i;
                hb += 1 << i;
            }
        }
        ha++;
        hb++;
        return {Ancestor[a][0], ha, hb};
    }
}

void hld(int u, int p, int c) {

    int v, sizeMax, k;

    Comp[u] = {c, CompVertecies[c].size()};
    CompVertecies[c].push_back(u);

    sizeMax = 0;
    k = -1;
    for (pair<int, int> e : Graph[u]) {
        v = e.first;
        if (v != p) {
            if (sizeMax < Size[v]) {
                sizeMax = Size[v];
                k = v;
            }
        }
    }

    for (pair<int, int> e : Graph[u]) {
        v = e.first;
        if (v != p) {
            if (v == k) {
                CompEdges[c].push_back(e.second);
                hld(v, u, c);
            }
            else {
                C++;
                CompVertecies[C - 1].push_back(u);
                CompEdges[C - 1].push_back(e.second);
                hld(v, u, C - 1);
            }
        }
    }
}

void apply(int u, int h) {

    int c, d, a, b;

    while (h != 0) {
        tie(c, d) = Comp[u];
        if (d >= h) {
            a = d - h;
            h = 0;
        }
        else {
            a = 0;
            h -= d;
            u = CompVertecies[c][0];
        }
        b = d - 1;
        Intervals[c].push_back({a, b});
        Jobs.push_back(c);
    }
}

int main() {

    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n, m, k, a, b, s, t, tNew, u, hu, ht, aNew, bNew, ans;
    ve<int> Solution;

    cin >> n >> m >> k;

    L = ceil(log2(n));
    Depth = ve<int>(n);
    Size = ve<int>(n);
    Comp = ve<pair<int, int>>(n);
    Graph = ve<ve<pair<int, int>>>(n);
    Ancestor = ve<ve<int>>(n, ve<int>(L));
    CompVertecies = ve<ve<int>>(n);
    CompEdges = ve<ve<int>>(n);
    Intervals = ve<ve<pair<int, int>>>(n);

    for (int i = 0; i < n - 1; i++) {
        cin >> a >> b;
        a--;
        b--;
        Graph[a].push_back({b, i});
        Graph[b].push_back({a, i});
    }

    dfs(0, -1, 0);

    C = 0;
    for (pair<int, int> e : Graph[0]) {
        C++;
        CompVertecies[C - 1].push_back(0);
        CompEdges[C - 1].push_back(e.second);
        hld(e.first, 0, C - 1);
    }

    SegTrees = vector<SegmentTree>(C);

    for (int i = 0; i < C; i++) {
        SegTrees[i].build(CompEdges[i].size());
    }

    for (int i = 0; i < m; i++) {
        cin >> s;
        cin >> t;
        t--;
        for (int j = 1; j < s; j++) {
            cin >> u;
            u--;
            tie(tNew, hu, ht) = lca(u, t);
            apply(u, hu);
            apply(t, ht);
            t = tNew;
        }
        for (int j : Jobs) {
            if (Intervals[j].size() > 0) {
                sort(Intervals[j].begin(), Intervals[j].end());
                tie(a, b) = Intervals[j][0];
                for (int l = 1; l < Intervals[j].size(); l++) {
                    tie(aNew, bNew) = Intervals[j][l];
                    if (aNew <= b) b = max(b, bNew);
                    else {
                        SegTrees[j].rangeUpdate(a, b, 1);
                        a = aNew;
                        b = bNew;
                    }
                }
                SegTrees[j].rangeUpdate(a, b, 1);
                Intervals[j].clear();
            }
        }
        Jobs.clear();
    }

    ans = 0;
    for (int i = 0; i < C; i++) {
        for (int j = 0; j < CompEdges[i].size(); j++) {
            if (SegTrees[i].query(j) >= k) {
                Solution.push_back(CompEdges[i][j]);
                ans++;
            }
        }
    }

    sort(Solution.begin(), Solution.end());

    cout << ans << endl;

    for (int i = 0; i < ans; i++) {
        cout << Solution[i] + 1 << " ";
    }
    cout << endl;

    return 0;
}

Compilation message (stderr)

railway.cpp: In function 'int main()':
railway.cpp:241:35: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
                 for (int l = 1; l < Intervals[j].size(); l++) {
                                 ~~^~~~~~~~~~~~~~~~~~~~~
railway.cpp:259:27: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
         for (int j = 0; j < CompEdges[i].size(); j++) {
                         ~~^~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 8 / 8
• Subtask 2: 15 / 15
• Subtask 3: 7 / 7
• Subtask 4: 29 / 29
• Subtask 5: 16 / 16
• Subtask 6: 25 / 25