Нахождение вершин в графе, в которые можно прийти по ребру с весом 1
Есть неориентированный взвешенный граф с весами рёбер 1 или 0, в графе n вершин и n рёбер. Мы стартуем из вершины К, нужно найти все вершины, последнее ребро пути до которых было с весом 1. Кроме того, нельзя посещать одно и то же ребро дважды (вершины посещать дважды можно) У меня вышел вот такой вот код:
#include <iostream>
#include <vector>
#include <algorithm>
#include <set>
typedef long long ll;
using namespace std;
bool used[1000005];
bool _used[1000005];
vector <vector <pair<ll, ll>>> g(1e6 + 5);
vector <ll> p(1e6 + 5);
set <ll> ans;
ll n;
ll node;
void dfs(ll v)
{
used[v] = true;
for (ll j = 0; j < g[v].size(); j++) {
if (used[g[v][j].first] == true && p[v] != g[v][j].first) {
p[g[v][j].first] = v;
node = g[v][j].first;
}
else
if (used[g[v][j].first] == false) {
p[g[v][j].first] = v;
dfs(g[v][j].first);
}
}
}
void DFS(ll v, ll clr)
{
_used[v] = true;
for (ll j = 0; j < g[v].size(); j++) {
if (clr == 1) {
ans.insert(v);
}
if (_used[g[v][j].first] == false) {
DFS(g[v][j].first, g[v][j].second);
}
}
}
vector <ll> getCycle()
{
vector <ll> c;
ll frst = node;
for (ll i = 0; i < n; i++) {
node = p[node];
c.push_back(node);
if (node == frst) {
break;
}
}
reverse (c.begin (), c.end ());
return c;
}
signed main()
{
cout.tie(nullptr);
cin.tie(nullptr);
ios_base::sync_with_stdio(false);
cin >> n;
for (ll i = 0; i < n; i++) {
ll a, b, c;
cin >> a >> b >> c;
g[a].push_back({ b, c });
g[b].push_back({ a, c });
}
ll k;
cin >> k;
for (int i = 1; i <= n; ++i)
if (used[i] == false)
dfs(i);
vector <ll> inf = getCycle();
DFS(k, 0);
inf.push_back(inf[0]);
for (int i = 0; i < inf.size() - 1; i++) {
ll a = inf[i];
ll b = inf[i + 1];
for (int j = 0; j < g[a].size(); j++) {
if (g[a][j].first == b) {
if (g[a][j].second == 1) {
ans.insert(a);
ans.insert(b);
}
break;
}
}
}
/*
cout << '\n';
auto _it = ans.begin();
for (ll i = 0; i < inf.size(); i++) {
cout << inf[i] << ' ';
_it++;
}
cout << '\n';
*/
ans.erase(k);
cout << ans.size();
}
Т.к. Наш граф - это дерево с дополнительным ребром. Если бы граф был обычным деревом, то обычный DFS дал бы нам ответ на этот вопрос. Но т.к. у графа есть ещё одно ребро, то в графе существует цикл. Причём, если в этом цикле есть ребро с весом 1, то обе вершины входят в ответ (т.к. это цикл мы можем идти либо по часовой стрелке либо против часовой). Вроде бы это решение звучит логично, но оно неверно т.к. я получаю Wrong Answer на 7-ом тесте. Я правда не знаю что тут может быть не так, я час вручную проверял все графы, но ответ всегда был верным.
Вот так он задаётся:
7
7 6 1
6 2 0
2 1 0
2 3 1
1 4 0
4 3 1
3 5 0
2
