Submission #1199404

#	Time	Username	Problem	Language	Result	Execution time	Memory
1199404		m_bezrutchka	Railway (BOI17_railway)	C++20	100 / 100	355 ms	211468 KiB

railway

// submetendo minha solucao de uns anos atras que tava no CSES
// small-to-large, mas ta mais complicada do que precisava ser
#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
 
const int MAXN = (int)1e6 + 10;
 
vector <pii> graph[MAXN];
vector <int> queries_in_node[MAXN];
int query_sz[MAXN];
int n, m, K;
 
int parent[MAXN];
int deg[MAXN];
int parent_idx[MAXN];
bool seen[MAXN];
vector <int> leaves;
 
int set_id[MAXN];
set <int> existing[MAXN];
set <int> completed[MAXN];
map <int, int> query_to_cnt[MAXN];
queue <int> bfs_queue;
 
int answer[MAXN];
vector <int> ans;
 
void dfs(int v) {
  seen[v] = true;
  for (int k = 0; k < (int) graph[v].size(); k++) {
    int w = graph[v][k].first;
    int idx = graph[v][k].second;
    if (seen[w]) {
      continue;
    }
    parent[w] = v;
    parent_idx[w] = idx;
    dfs(w);
    deg[v]++;
  }
  if (deg[v] == 0) {
    leaves.push_back(v);
  }
}
 
void add_to_query(int v, int q_idx, int val) {
  // printf("add_to_query(%d, %d, %d)\n", v, q_idx, val);
  int cur_cnt;
  if (query_to_cnt[v].find(q_idx) == query_to_cnt[v].end()) {
    cur_cnt = val;
  } else {
    cur_cnt = query_to_cnt[v][q_idx] + val;
  }
  query_to_cnt[v][q_idx] = cur_cnt;
  if (existing[v].find(q_idx) == existing[v].end()) {
    existing[v].insert(q_idx);
  }
  if (query_to_cnt[v][q_idx] == query_sz[q_idx]) {
    if (completed[v].find(q_idx) == completed[v].end()) {
      completed[v].insert(q_idx);
    }
  }
}
 
int get_answer(int v) {
  v = set_id[v];
  int cur_ans = (int) existing[v].size() - (int) completed[v].size();
  return cur_ans;
}
 
void init_node(int v) {
  set_id[v] = v;
  for (int k = 0; k < (int) queries_in_node[v].size(); k++) {
    int qi = queries_in_node[v][k];
    add_to_query(v, qi, 1);
  }
}
 
void merge_nodes(int v, int w) {
  // printf("merge_nodes(%d, %d)\n", v, w);
  int a = set_id[v], b = set_id[w];
  // printf("a = set_id[%d] = %d; b = set_id[%d] = %d\n", v, a, w, b);
  int sz_a = (int) existing[a].size();
  int sz_b = (int) existing[b].size();
  // printf("sz_a = %d, sz_b = %d\n", sz_a, sz_b);
  if (sz_a < sz_b) {
    swap(a, b);
    // printf("swapped, now a = %d and b = %d\n", a, b);
  } else {
    // printf("did not swap\n");
  }
 
  map<int, int>::iterator it = query_to_cnt[b].begin();
  while (it != query_to_cnt[b].end()) {
    int idx = it->first;
    int cnt = it->second;
    add_to_query(a, idx, cnt);
    it++;
  }
  query_to_cnt[b].clear();
  completed[b].clear();
  existing[b].clear();
  // printf("cleared vals for %d\n", b);
 
  if (sz_a < sz_b) {
    swap(set_id[v], set_id[w]);
    // printf("swapped set_ids, now set_id[%d] = %d\n", v, set_id[v]);
  }
}
 
void bfs() {
  for (int i = 0; i < (int) leaves.size(); i++) {
    int v = leaves[i];
    init_node(v);
    int idx = parent_idx[v];
    if (idx != 0) {
      answer[idx] = get_answer(v);
    }
    int w = parent[v];
    if (w != 0) {
      deg[w]--;
      if (deg[w] == 0) {
        bfs_queue.push(w);
      }
    }
  }
 
  while (!bfs_queue.empty()) {
    int v = bfs_queue.front();
    bfs_queue.pop();
    init_node(v);
    for (int i = 0; i < (int) graph[v].size(); i++) {
      int w = graph[v][i].first;
      if (w == parent[v]) {
        continue;
      }
      merge_nodes(v, w);
    }
    int idx = parent_idx[v];
    if (idx != 0) {
      answer[idx] = get_answer(v);
    }
    int w = parent[v];
    if (w != 0) {
      deg[w]--;
      if (deg[w] == 0) {
        bfs_queue.push(w);
      }
    }
  }
}
 
int main() {
  scanf("%d %d %d", &n, &m, &K);
  for (int i = 1; i < n; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    graph[a].push_back(make_pair(b, i));
    graph[b].push_back(make_pair(a, i));
  }
  dfs(1);
 
  for (int i = 1; i <= m; i++) {
    int si, x;
    scanf("%d", &query_sz[i]);
    for (int k = 0; k < query_sz[i]; k++) {
      scanf("%d", &x);
      queries_in_node[x].push_back(i);
    }
  }
 
  bfs();
  for (int i = 1; i < n; i++) {
    if (answer[i] >= K) {
      ans.push_back(i);
    }
  }
 
  printf("%d\n", (int) ans.size());
  for (int i = 0; i < (int) ans.size(); i++) {
    printf("%d%c", ans[i], (i == ((int) ans.size()) - 1) ? '\n' : ' ');
  }
  return 0;
}

Compilation message (stderr)

railway.cpp: In function 'int main()':
railway.cpp:157:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  157 |   scanf("%d %d %d", &n, &m, &K);
      |   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
railway.cpp:160:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  160 |     scanf("%d %d", &a, &b);
      |     ~~~~~^~~~~~~~~~~~~~~~~
railway.cpp:168:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  168 |     scanf("%d", &query_sz[i]);
      |     ~~~~~^~~~~~~~~~~~~~~~~~~~
railway.cpp:170:12: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  170 |       scanf("%d", &x);
      |       ~~~~~^~~~~~~~~~

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