Interference - BIO 2024 Round 2

/**
 * @file interference.cpp
 * @version 4.0 (of problem statement)
 * @date 2024-04-01
 *
 * usage:
 *      Read from / write to default input.txt and output.txt
 *          interference.exe
 *      Read from / write to custom files:
 *          interference.exe in.txt out.txt
 */
#include <iostream>
#include <vector>
#include <functional>
#include <cassert>

using namespace std;

void fileIO(int argc, char *argv[]);

int main(int argc, char *argv[]) {
    fileIO(argc, argv); // remove for standard input/output

    int s, d; cin >> s >> d;
    vector<int> f(1+s);
    // set f to 3 for all nodes in disrepair
    for (int i = 0; i < d; ++i) {
        int disrepair; cin >> disrepair;
        f[disrepair] = 3;
    }

    vector<vector<pair<int, bool>>> g(1+s);
    for (int u, v, w; cin >> u >> v >> w && u != -1; ) {
        g[u].emplace_back(v, w);
        g[v].emplace_back(u, w);
    }

    function<void(int, int)> dfsSetup = [&](int node, int parent = 0) {
        f[node] = 1;
        for (auto [child, far] : g[node]) {
            if (child != parent) {
                if (f[child]) {
                    assert(f[child] == 3);
                    if (!far) f[node] = 2;
                } else {
                    dfsSetup(child, node);
                    assert(f[child] < 3); // roots are always assigned values < 3, we will flip the subtrees later if we want this to be a far edge
                }
            }
        }
    };
    function<void(int, int, bool)> dfsCorrect = [&](int node, int parent = 0, bool flipSubtree = false) {
        for (auto [child, far] : g[node]) {
            if (child != parent) {
                assert(f[child]);
                if (f[child] != 3) dfsCorrect(child, node, flipSubtree ^ far);
            }
        }
        if (flipSubtree) f[node] = 6 - f[node];
    };

    for (int i = 1; i <= s; ++i) {
        if (!f[i]) {
            dfsSetup(i, 0);
            dfsCorrect(i, 0, false);
        }
        cout << f[i] << '\n';
    }
    return 0;
}

void fileIO(int argc, char *argv[]) {
    const char * in = "input.txt", * out = "output.txt";
    if (argc > 1) in = argv[1];
    if (argc > 2) out = argv[2];
    freopen(in, "r", stdin);
    freopen(out, "w", stdout);
}