Submission #1043530

#	Time	Username	Problem	Language	Result	Execution time	Memory
1043530		ssense	Railway (BOI17_railway)	C++14	100 / 100	238 ms	63572 KiB

railway

#include <iostream>
#include <algorithm>
#include <vector>
#include <array>
#include <set>
#include <map>
#include <queue>
#include <stack>
#include <list>
#include <chrono>
#include <random>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <cstring>
#include <iomanip>
#include <bitset>
#include <limits.h>
#include <cassert>

#define fr first
#define sc second

using namespace std;

#define int long long

const int N = (1e5)+5;

int n, m, k;

vector<int> adj[N];

int par[N][22], depth[N];

void dfs(int u, int p)
{
    depth[u] = depth[p]+1;
    for(auto v : adj[u])
    {
        if(v != p)
        {
            par[v][0] = u;
            dfs(v, u);
        }
    }
}

int jump(int a, int len)
{
    // cout << a << " " << len << endl;
    for(int i = 0; i <= 20; i++)
    {
        if((len&(1<<i)) != 0)
        {
            a = par[a][i];
        }
    }
    // cout << a << endl;
    return a;
}

int lca(int a, int b)
{
    if(depth[a] < depth[b])
    {
        swap(a, b);
    }
    a = jump(a, depth[a]-depth[b]);
    if(a == b){return a;}
    for(int i = 20; i >= 0; i--)
    {
        if(par[a][i] != par[b][i])
        {
            a = par[a][i];
            b = par[b][i];
        }
    }
    return par[a][0];
}

set<pair<int, int> > tree[N];

map<pair<int, int>, int> mpe;

vector<int> ans;

void dfs2(int u, int p)
{
    for(auto v : adj[u])
    {
        if(v != p)
        {
            dfs2(v, u);
        }
    }
    for(auto v : adj[u])
    {
        if(v != p)
        {
            if(tree[v].size() > tree[u].size())
            {
                swap(tree[v], tree[u]);
            }
        }
    }
    for(auto v : adj[u])
    {
        if(v != p)
        {
            tree[u].insert(tree[v].begin(), tree[v].end());
        }
    }
    auto lower = tree[u].lower_bound(std::make_pair(u, INT_MIN));
    auto upper = tree[u].upper_bound(std::make_pair(u, INT_MAX));
    tree[u].erase(lower, upper);
    // cout << "NODE " << u << ": ";
    // for(auto x: tree[u])
    // {
    //     cout << " (" <<x.fr << ", " << x.sc << ") ";
    // }
    // cout << endl;
    if(tree[u].size() >= k)
    {
        ans.push_back(mpe[make_pair(min(u, p), max(u, p))]);
    }
}

void solve()
{
    cin >> n >> m >> k;
    for(int i = 0; i < n-1; i++)
    {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
        mpe[make_pair(min(u, v), max(u, v))] = i+1;
    }
    dfs(1, 0);
    for(int i = 1; i <= 20; i++)
    {
        for(int j = 1; j <= n; j++)
        {
            par[j][i] = par[par[j][i-1]][i-1];
        }
    }
    // for(int i = 1; i <= n; i++)
    // {
    //     cout << i << " " << depth[i] << endl;
    // }
    for(int i = 0; i < m; i++)
    {
        int cnt;
        cin >> cnt;
        vector<int> b(cnt);
        int lc;
        for(int j = 0; j < cnt; j++)
        {
            cin >> b[j];
            if(j == 0){lc = b[0];}
            else{lc = lca(lc, b[j]);}
        }
        assert(lc != 0);
        // cout << "PATANUL " << i << " LCA " << lc << endl;
        for(int j = 0; j < cnt; j++)
        {
            tree[b[j]].insert(make_pair(lc, i));
        }
    }
    dfs2(1, 0);
    //assert(tree[1].size() == 0);
    cout << ans.size() << endl;
    sort(ans.begin(), ans.end());
    for(auto x : ans)
    {
        cout << x << " ";
    }
}

int32_t main()
{
    int t = 1;
    //cin >> t;
    while(t--)
    {
        solve();
    }
}

/*
6 3 2
1 3
2 3
3 4
6 4
4 5
4 1 3 2 5
2 6 3
2 3 2
*/

/*
#include <iostream>
#include <algorithm>
#include <vector>
#include <array>
#include <set>
#include <map>
#include <queue>
#include <stack>
#include <list>
#include <chrono>
#include <random>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <cstring>
#include <iomanip>
#include <bitset>
#include <cassert>
*/

Compilation message (stderr)

railway.cpp: In function 'void dfs2(long long int, long long int)':
railway.cpp:123:23: warning: comparison of integer expressions of different signedness: 'std::set<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
  123 |     if(tree[u].size() >= k)
      |        ~~~~~~~~~~~~~~~^~~~

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