Нахождение вершин в графе, в которые можно прийти по ребру с весом 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


Ответы (0 шт):