#include <bits/stdc++.h>
using namespace std;
#define ms(x, a) memset(x, a, sizeof x)
using ll = long long;
const int mod = 1e9 + 7, inf = 0x3f3f3f3f;
const int MN = 0;
// #ifndef ONLINE_JUDGE
// int count_common_roads(vector<int> &r);
// #endif
int N, M;
vector<int> A, B;
vector<int> royal;
vector<int> span;
vector<bool> vis;
map<pair<int, int>, int> id_lookup;
struct DisjointSet{
vector<int> p, sz;
DisjointSet(int n){
p.resize(n);
iota(p.begin(), p.end(), 0);
sz.resize(n, 1);
}
int Find(int x){ return p[x] == x ? x : p[x] = Find(p[x]); }
bool Union(int a, int b){
a = Find(a), b = Find(b);
if (a == b) return 0;
if (sz[b] > sz[a]) swap(a, b);
p[b] = a; sz[a] += sz[b];
return 1;
}
};
int complete_forest(vector<int> &e){
DisjointSet dsu(N);
for (int i : e) dsu.Union(A[i], B[i]);
int sub = 0;
for (int i : span){
if (dsu.Union(A[i], B[i])){
if (royal[i] == 1) sub++;
e.push_back(i);
}
}
return sub;
}
int count_common_roads_forest(vector<int> &e){
int sub = complete_forest(e);
return count_common_roads(e) - sub;
}
void dfs(int v){
vis[v] = 1;
for (int i = 0; i < M; ++i){
if (A[i] != v && B[i] != v) continue;
int to = A[i] ^ B[i] ^ v;
if (vis[to]) continue;
span.push_back(i);
dfs(to);
}
}
void build_span(){
dfs(0);
for (int i = 0; i < N - 1; ++i){
vector<int> e;
int a = A[span[i]], b = B[span[i]];
int c, cur, mn = inf, mx = 0;
if (a == 0 || b == 0){
if (a == 1 || b == 1) c = 2;
else c = 1;
}
else c = 0;
e = {id_lookup[{a, b}], id_lookup[{a, c}]};
cur = count_common_roads_forest(e);
mn = min(mn, cur); mx = max(mx, cur);
e = {id_lookup[{c, b}], id_lookup[{a, b}]};
cur = count_common_roads_forest(e);
mn = min(mn, cur); mx = max(mx, cur);
e = {id_lookup[{a, c}], id_lookup[{c, b}]};
cur = count_common_roads_forest(e);
mn = min(mn, cur); mx = max(mx, cur);
if (mn == mx) royal[span[i]] = -1;
else{
if (cur == mn) royal[span[i]] = 1;
else royal[span[i]] = -1;
}
// debug spanning tree
// cout << A[span[i]] << " " << B[span[i]] << " " << royal[span[i]] << '\n';
// cout << mn << " " << cur << " " << mx << '\n';
}
}
vector<int> find_degrees(){
vector<int> degree(N);
for (int i = 0; i < N; ++i){
vector<int> e;
for (int j = 0; j < M; ++j){
if (A[j] == i || B[j] == i){
e.push_back(j);
}
}
degree[i] = count_common_roads_forest(e);
}
return degree;
}
int solve(vector<int> &e, int l, int r){
if (l + 1 == r) return e[l];
int mid = (l + r) >> 1;
vector<int> left_edges;
for (int i = l; i < mid; ++i) left_edges.push_back(e[i]);
if (count_common_roads_forest(left_edges) == 1) return solve(e, l, mid);
else return solve(e, mid, r);
}
vector<int> find_roads(int n, vector<int> a, vector<int> b){
N = n; M = a.size();
A = a, B = b;
for (int i = 0; i < M; ++i){
id_lookup[{A[i], B[i]}] = i;
id_lookup[{B[i], A[i]}] = i;
}
royal.resize(M);
vis.resize(N);
build_span();
vector<int> degree = find_degrees();
// debug degrees
// for (int i = 0; i < N; ++i) cout << degree[i] << (i == N - 1 ? '\n' : ' ');
// debug royal edges in span tree
// for (int i = 0; i < M; ++i) cout << royal[i] << " ";
// cout << '\n';
queue<int> q;
for (int i = 0; i < N; ++i){
if (degree[i] == 1) q.push(i);
}
vis = vector<bool>(N, 0);
while (!q.empty()){
if (q.size() == 1) break;
int v = q.front(); q.pop();
vis[v] = 1;
vector<int> e;
for (int i = 0; i < M; ++i){
if (A[i] != v && B[i] != v) continue;
int to = A[i] ^ B[i] ^ v;
if (vis[to]) continue;
e.push_back(i);
}
int good_id = solve(e, 0, e.size());
for (int i = 0; i < e.size(); ++i){
if (e[i] == good_id){
royal[e[i]] = 1;
int to = A[e[i]] ^ B[e[i]] ^ v;
degree[to]--;
if (degree[to] == 1) q.push(to);
}
else royal[e[i]] = -1;
}
}
// debug royal edges
// for (int i = 0; i < M; ++i) cout << royal[i] << " ";
// cout << '\n';
vector<int> ans;
for (int i = 0; i < M; ++i){
if (royal[i] == 1) ans.push_back(i);
}
return ans;
}
// #ifndef ONLINE_JUDGE
// vector<bool> is_royal_ans;
// int total_queries = 0;
// int count_common_roads(vector<int> &r){
// total_queries++;
// int common = 0;
// for(int i = 0; i < N - 1; i++) {
// bool is_common = is_royal_ans[r[i]];
// if (is_common)
// common++;
// }
// return common;
// }
// int main(){
// // cin.tie(0)->sync_with_stdio(0);
// int n; n = 4;
// vector<int> a = {0, 0, 0, 1, 1, 2};
// vector<int> b = {1, 2, 3, 2, 3, 3};
// is_royal_ans.resize(a.size());
// is_royal_ans[0] = is_royal_ans[1] = is_royal_ans[5] = 1;
// vector<int> ans = find_roads(n, a, b);
// for (int i : ans){
// cout << i << " ";
// }
// cout << '\n';
// cout << "Queries: " << total_queries << '\n';
// }
// #endif
Compilation message
simurgh.cpp: In function 'int count_common_roads_forest(std::vector<int>&)':
simurgh.cpp:51:12: error: 'count_common_roads' was not declared in this scope; did you mean 'count_common_roads_forest'?
51 | return count_common_roads(e) - sub;
| ^~~~~~~~~~~~~~~~~~
| count_common_roads_forest
simurgh.cpp: In function 'std::vector<int> find_roads(int, std::vector<int>, std::vector<int>)':
simurgh.cpp:163:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
163 | for (int i = 0; i < e.size(); ++i){
| ~~^~~~~~~~~~
