Submission #1110838

#	Time	Username	Problem	Language	Result	Execution time	Memory
1110838		Icelast	Railway (BOI17_railway)	C++17	100 / 100	142 ms	59004 KiB

railway

#include <iostream>
#include <bits/stdc++.h>
#define ll long long
using namespace std;
const ll maxn = 2*1e5+5, INF = 4e18+9;
struct Tree{
    int n, root;
    vector<int> depth, head, sz, pa;
    vector<int> f;
    Tree(int n, int root, vector<vector<int>> &adj): n(n), root(root){
        depth.resize(n+1, -1);
        head.resize(n+1);
        sz.resize(n+1, 0);
        pa.resize(n+1, -1);
        auto rootTree = [&](auto rootTree, int u, int p) -> void{
            depth[u] = depth[p]+1;
            pa[u] = p;
            sz[u] = 1;
            for(int v : adj[u]){
                if(v == p) continue;
                rootTree(rootTree, v, u);
                sz[u] += sz[v];
            }
        };
        rootTree(rootTree, root, 0);
        auto decompose = [&](auto decompose, int u, int h) -> void{
            head[u] = h;
            int heavy = -1;
            for(int v : adj[u]){
                if(v == pa[u]) continue;
                if(heavy == -1 || sz[heavy] < sz[v]) heavy = v;
            }
            if(heavy != -1) decompose(decompose, heavy, h);
            for(int v : adj[u]){
                if(v == pa[u] || v == heavy) continue;
                decompose(decompose, v, v);
            }
        };
        decompose(decompose, root, root);
        f = [&]{
            vector<array<int, 2>> first(n+1);
            vector<array<int, 2>> second(n+1);
            function<void(int, int)> dfs0 = [&](int u, int p) {
                first[u] = second[u] = {0, -1};
                for (int v : adj[u]) {
                    if (v == p) {
                        continue;
                    }
                    dfs0(v, u);
                    auto fetch = first[v];
                    fetch[0] += 1;
                    fetch[1] = v;
                    if (fetch > first[u]) {
                        swap(fetch, first[u]);
                    }
                    if (fetch > second[u]) {
                        swap(fetch, second[u]);
                    }
                }
            };
            dfs0(1, 0);

            function<void(int, int)> dfs = [&](int u, int p) {
                for (int v : adj[u]) {
                    if (v == p) {
                        continue;
                    }
                    auto fetch = first[u][1] == v ? second[u] : first[u];
                    fetch[0] += 1;
                    fetch[1] = u;
                    if (fetch > first[v]) {
                        swap(fetch, first[v]);
                    }
                    if (fetch > second[v]) {
                        swap(fetch, second[v]);
                    }
                    dfs(v, u);
                }
            };
            dfs(1, 0);
            vector<int> f(n+1);
            for (int u = 1; u <= n; u++) {
                f[u] = first[u][0];
            }
            return f;
        }();
    };
    int lca(int u, int v){
        for(; head[u] != head[v]; v = pa[head[v]]){
            if(depth[head[u]] > depth[head[v]]) swap(u, v);
        }
        if(depth[u] > depth[v]) swap(u, v);
        return u;
    }
    int dist(int u, int v) {
        return depth[u]+depth[v]-2*depth[lca(u, v)];
    }
};
template <class T>
struct Fenwick {
    int n, log;
    vector<T> bit;

    Fenwick(int n) : n(n), log(32 - __builtin_clz(n + 1)), bit(n + 1, 0) {}

    void add(int i, T delta) {
        for (; i <= n; i += i & -i) {
            bit[i] += delta;
        }
    }

    T sum(int i) {
        T res = 0;
        for (; i > 0; i -= i & -i) {
            res += bit[i];
        }
        return res;
    }
    T sum(int l, int r) {
        T res = 0;
        return sum(r)-sum(l-1);
    }
    int kth(T k) {
        T sum = 0;
        int pos = 0;
        for (int l = log - 1; l >= 0; l--) {
            if (pos + (1 << l) <= n && sum + bit[pos + (1 << l)] <= k) {
                pos += 1 << l;
                sum += bit[pos];
            }
        }
        return pos;
    }
};
void solve(){
    int n, m, k;
    cin >> n >> m >> k;
    vector<vector<int>> adj(n+1);
    vector<vector<pair<int, int>>> adj2(n+1);
    for(int i = 1; i < n; i++){
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
        adj2[u].push_back({v, i});
        adj2[v].push_back({u, i});
    }
    int timer = 0;
    vector<int> lt(n+1), rt(n+1);
    vector<int> e(n+1);
    auto dfs = [&](auto dfs, int u, int p) -> void{
        timer++;
        lt[u] = timer;
        for(auto it : adj2[u]){
            int v = it.first, id = it.second;
            if(v == p) continue;
            e[v] = id;
            dfs(dfs, v, u);
        }
        rt[u] = timer;
    };
    dfs(dfs, 1, 0);
    Tree Y(n, 1, adj);
    Fenwick<int> bit(n+1);
    auto addPath = [&](int u, int v) -> void{
        bit.add(lt[u], 1);
        bit.add(lt[v], 1);
        bit.add(lt[Y.lca(u, v)], -2);
    };
    auto get = [&](int v) -> int{
        return bit.sum(lt[v], rt[v]);
    };
    for(int i = 1; i <= m; i++){
        int s;
        cin >> s;
        vector<int> c(s+1);
        for(int j = 1; j <= s; j++){
            cin >> c[j];
        }
        sort(c.begin()+1, c.end(), [&](int a, int b){return lt[a] < lt[b];});
        c.push_back(c[1]);
        for(int i = 1; i <= s; i++){
            addPath(c[i], c[i+1]);
        }
    }
    vector<int> ans;
    for(int i = 2; i <= n; i++){
        if(get(i) >= 2*k){
            ans.push_back(e[i]);
        }
    }
    sort(ans.begin(), ans.end());
    cout << ans.size() << "\n";
    for(int i : ans) cout << i << " ";
}
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    solve();
}

Compilation message (stderr)

railway.cpp: In instantiation of 'T Fenwick<T>::sum(int, int) [with T = int]':
railway.cpp:171:36:   required from here
railway.cpp:120:11: warning: unused variable 'res' [-Wunused-variable]
  120 |         T res = 0;
      |           ^~~

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