Submission #245787

#	Time	Username	Problem	Language	Result	Execution time	Memory
245787		santaclaus03	Railway (BOI17_railway)	C++14	0 / 100	632 ms	41628 KiB

railway

#include <bits/stdc++.h>
#define INF 1000000000
#define loop(var, start, cond, step) for (int var = start; var cond; var += step)
#define range(var, start, endExclusive) loop(var, start, < endExclusive, 1)
#define repeat(times) range(____i, 0, times)
#define foreach(idx, element, arr) for (int i = 0, element = arr[0]; i < (int) arr.size(); element = arr[++i])
#define dbg(expr) cout << #expr << " = " << expr << endl;
using namespace std;

using vi = vector<int>;
using ii = pair<int, int>;
using vvi = vector<vi>;
using vii = vector<ii>;
using ll = long long;

template <typename T>
istream& operator>>(istream& is, vector<T>& v) {
   range(i, 0, v.size()) cin >> v[i]; 
   return is;
}

template <typename T>
bool in_range(T lo, T hi, T t) {
    return !(t < lo) && t < hi;
}

int n;
vi s, modified;
vvi tree, anc;
vi d, f, pref;
vii edges;
map<ii, int> cnt;
int t = 1;

void update(int j, int delta, int a = 1, int b = n * 2, int i = 1) {
    if (a == b) {
        s[i] += delta;
        modified.push_back(i);
    }
    else {
        int m = (a + b) / 2;
        if (j <= m) update(j, delta, a, m, i * 2);
        else update(j, delta, m + 1, b, i * 2 + 1);
        s[i] = s[i * 2] + s[i * 2 + 1];
        //cout << a << ".." << b << " = " << s[i] << endl;
        modified.push_back(i);
    }
} 

int query(int l, int r, int a = 1, int b = n * 2, int i = 1) {
    if (r < a || b < l) return 0;
    if (l <= a && b <= r) return s[i];
    int m = (a + b) / 2;
    int left = query(l, r, a, m, i * 2);
    int right = query(l, r, m + 1, b, i * 2 + 1);
    //cout << "l=" << l << ", r=" << r << ", a=" << a << ", b = "<<  b << ", sum = " << left + right << endl;
    return left + right;
}

void pre(int u, int p) {
    //d[u] = t++; cout << "discovery of " << u + 1 << " is " << d[u] << endl;
    anc[u][0] = p;
    for (int i = 1; i < 21; ++i) anc[u][i] = anc[anc[u][i - 1]][i-1];
    for (int v : tree[u]) if (v != p) pre(v, u);
    f[u] = t++;
    //cout << "finish of " << u + 1 << " is " << f[u] << endl;
}

bool is_anc(int p, int a) {
    return d[p] <= d[a] && f[a] <= f[p]; 
}

int lca(int a, int b) {
    if (is_anc(a, b)) return a;
    if (is_anc(b, a)) return b;
    for (int step = 19; step >= 0; --step) {
        if (!is_anc(anc[a][step], b)) a = anc[a][step];
    }
    return anc[a][0];
}

int dfs(int u, int p) {
    int res = pref[u];
    for (int v : tree[u]) if (v != p) {
        int n = dfs(v, u);
        res += n;
        cnt[{ min(u, v), max(u, v) }] = n;
        //cout << "edge " << u + 1 << ", " << v + 1 << " needed " << n << " times\n";
    }
    return res;
}

int main() {
    #ifdef DEBUG
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
    #endif
    int m, k; cin >> n >> m >> k;
    tree.resize(n), anc.resize(n, vi(21)), d.resize(n), 
    f.resize(n), pref.resize(n), edges.reserve(n - 1);
    s.resize(n * 8, 0);
    repeat(n - 1) {
        int a, b; cin >> a >> b; a--; b--;
        tree[a].push_back(b), tree[b].push_back(a);
        edges.emplace_back(min(a, b), max(a, b));
    }
    pre(0, 0);
    //for (int i = 0; i < n; ++i) cout << "parent of " << i + 1 << " is " << anc[i][0] + 1<< endl;
    repeat(m) {
        for (int i : modified) s[i] = 0;
        modified.clear();
        int c; cin >> c;
        vi nodes(c);
        for (int &n : nodes) {
            cin >> n; --n;
            pref[n]++;
            //cout << "update " << d[n] << "(" << n + 1 <<  ") by +1\n"; 
            update(d[n], +1);
        }
        sort(nodes.begin(), nodes.end(), [](int a, int b) { return is_anc(b, a); });
        int ca = nodes[0];
        for (int step = 19; step >= 0; --step) {
            if (query(d[anc[ca][step]], f[anc[ca][step]]) < c) ca = anc[ca][step]; 
        }
        ca = anc[ca][0];
        //cout << "common parent is " << ca + 1 << endl;
        for (int n : nodes) {
            //cout << "computing first LCA from " << n + 1 << endl;
            for (int step = 19; step >= 0; --step) {
                int a = anc[n][step];
                int l = d[a], r = f[a];
                int children = query(l, r);
                if (children < 2) n = a;
            }
            n = anc[n][0];
            //cout << "is " << n + 1 << endl;
            int other_children = query(d[n], f[n]) - 1; 
            //cout << other_children << " other children (" << d[n] << ".." << f[n] << ")\n";
            if (other_children == 0) continue;
            //cout << "update " << d[n] << " by -" << other_children << endl;
            update(d[n], -other_children);
            pref[n] -= other_children;
        }
        //cout << "update " << d[ca] << " by -1\n";
        update(d[ca], -1);
        pref[ca]--;
    }
    //for (int i = 0; i < n; ++i) cout << "v " << i + 1 << " = " << pref[i] << endl;
    dfs(0, 0);
    vi ans;
    for (int i = 0; i < edges.size(); ++i) {
        if (cnt[edges[i]] >= k) ans.push_back(i);
    }
    cout << ans.size() << endl;
    for (int e : ans) cout << e + 1 << " ";
    return 0;
}

Compilation message (stderr)

railway.cpp: In function 'int main()':
railway.cpp:151:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < edges.size(); ++i) {
                     ~~^~~~~~~~~~~~~~

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