Submission #797551

#	Time	Username	Problem	Language	Result	Execution time	Memory
797551		Tigerpants	Railway (BOI17_railway)	C++17	100 / 100	247 ms	61008 KiB

railway

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <numeric>
#include <functional>

using namespace std;

typedef long long int ll;
typedef vector<ll> vll;
typedef pair<ll, ll> pll;
typedef vector<vll> vvll;
typedef vector<pll> vpll;
typedef vector<vpll> vvpll;
typedef set<ll> sll;
typedef map<ll, sll> msll;

const ll MAXLOGN = 18;
const ll po2[MAXLOGN] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072};

#define rep(i, a, b) for (ll i = a; i < b; i++)
#define mp(a, b) make_pair(a, b)
#define sz(a) a.size()
#define pb(a) push_back(a)

ll n, m, k;
vvpll g;

vvll up;
vll edge;
vll depth;

vvpll charge;

void root_tree(ll cur, ll par) {
    up[cur][0] = par;
    for (vpll::iterator e = g[cur].begin(); e != g[cur].end(); e++) {
        if (e->first == par) {
            edge[cur] = e->second;
            depth[cur] = depth[par] + 1;
        } else {
            root_tree(e->first, cur);
        }
    }
}

ll lca(ll a, ll b) {
    // clamp to same depth
    if (depth[a] > depth[b]) swap(a, b);
    for (ll d = MAXLOGN - 1; d >= 0; d--) {
        if (depth[b] - po2[d] >= depth[a]) b = up[b][d];
    }
    
    if (a == b) return a;

    for (ll d = MAXLOGN - 1; d >= 0; d--) {
        if (up[a][d] != up[b][d]) {
            a = up[a][d];
            b = up[b][d];
        }
    }
    return up[a][0];
}

void union_merge_sll(sll& a, sll& b) {
    if (sz(b) > sz(a)) swap(a, b);
    for (sll::iterator itr = b.begin(); itr != b.end(); itr++) {
        a.insert(*itr);
    }
}

void union_merge_msll(msll& a, msll& b) {
    if (sz(b) > sz(a)) swap(a, b);
    for (msll::iterator itr = b.begin(); itr != b.end(); itr++) {
        union_merge_sll(a[itr->first], b[itr->first]);
    }
}

vll answer;
void pull_charge(ll cur, ll par, sll& ministers, msll& time) {
    for (vpll::iterator e = g[cur].begin(); e != g[cur].end(); e++) {
        if (e->first == par) continue;
        
        sll new_ministers;
        msll new_time;
        pull_charge(e->first, cur, new_ministers, new_time);
        union_merge_sll(ministers, new_ministers);
        union_merge_msll(time, new_time);
    }
    // add charge at node
    for (vpll::iterator itr = charge[cur].begin(); itr != charge[cur].end(); itr++) {
        ministers.insert(itr->first);
        time[itr->second].insert(itr->first);
    }

    // remove all with expired time == depth[cur]
    for (sll::iterator itr = time[depth[cur]].begin(); itr != time[depth[cur]].end(); itr++) {
        ministers.erase(*itr);
    }
    time.erase(depth[cur]);

    // decide whether we include edge in planning
    if (cur != 0) {
        if (sz(ministers) >= k) {
            answer.pb(edge[cur]);
        }
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    cin >> n >> m >> k;
    g.resize(n);
    rep(i, 1, n) {
        ll a, b;
        cin >> a >> b;
        a--; b--;
        g[a].pb(mp(b, i));
        g[b].pb(mp(a, i));
    }

    up.resize(n);
    rep(i, 0, n) up[i].resize(MAXLOGN);
    edge.resize(n);
    edge[0] = -1;
    depth.resize(n);
    rep(i, 0, n) depth[i] = 0;

    root_tree(0, -1);

    // prep lca
    rep(i, 1, MAXLOGN) {
        rep(j, 0, n) {
            if (up[j][i - 1] == -1) {
                up[j][i] = -1;
            } else {
                up[j][i] = up[up[j][i - 1]][i - 1];
            }
        }
    }

    charge.resize(n);

    // process queries
    rep(q, 0, m) {
        ll s;
        cin >> s;
        vll a(s);
        rep(i, 0, s) {
            cin >> a[i]; a[i]--;
        }

        ll lca_a = a[0];
        rep(i, 1, s) {
            lca_a = lca(lca_a, a[i]);
        }

        rep(i, 0, s) {
            charge[a[i]].pb(mp(q, depth[lca_a]));
        }
    }

    // calculate result
    sll ministers;
    msll time;
    pull_charge(0, -1, ministers, time);

    // answer
    sort(answer.begin(), answer.end());
    cout << sz(answer) << "\n";
    rep(i, 0, sz(answer)) {
        cout << answer[i] << " ";
    }
    cout << endl;

    return 0;
}

Compilation message (stderr)

railway.cpp: In function 'void pull_charge(ll, ll, sll&, msll&)':
railway.cpp:106:27: warning: comparison of integer expressions of different signedness: 'std::set<long long int>::size_type' {aka 'long unsigned int'} and 'll' {aka 'long long int'} [-Wsign-compare]
  106 |         if (sz(ministers) >= k) {
      |                           ^
railway.cpp: In function 'int main()':
railway.cpp:23:39: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   23 | #define rep(i, a, b) for (ll i = a; i < b; i++)
      |                                       ^
railway.cpp:175:5: note: in expansion of macro 'rep'
  175 |     rep(i, 0, sz(answer)) {
      |     ^~~

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