IOI '20 - Supertrees

#include "supertrees.h"
#include <vector>

using namespace std;

int construct(vector<vector<int>> p) {
    int n = p.size();
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            if (p[i][j] == 3) return 0;
        }
    }

    // Split into components
    vector<vector<int>> components;
    vector<int> myCmp(n, -1);
    for (int i = 0; i < n; ++i) {
        if (myCmp[i] != -1) continue;
        components.push_back({});
        for (int j = 0; j < n; ++j) {
            if (p[i][j]) {
                if (myCmp[j] != -1) return 0;
                components.back().push_back(j);
                myCmp[j] = (int)components.size()-1;
            }
        }
    }

    // Check components are separate
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            if ((myCmp[i] == myCmp[j]) ^ (p[i][j] != 0)) return 0;
        }
    }

    // Build answer here
    vector<vector<int>> edge(n, vector<int>(n));

    // Now solve separately for each component
    for (vector<int> &c : components) {
        // Split into 1-chains
        vector<vector<int>> chains;
        vector<int> myChain(n, -1);
        for (int i : c) {
            if (myChain[i] != -1) continue;
            chains.push_back({});
            for (int j : c) {
                if (p[i][j] == 1) {
                    if (myChain[j] != -1) return 0;
                    chains.back().push_back(j);
                    myChain[j] = (int)chains.size()-1;
                }
            }
        }

        // Check 1-chains are separate
        for (int i : c) {
            for (int j : c) {
                if ((myChain[i] == myChain[j]) ^ (p[i][j] == 1)) return 0;
            }
        }

        // Assign a chain leader to each chain
        int m = chains.size();
        vector<int> chainLeader(m);
        for (int i = 0; i < m; ++i) chainLeader[i] = chains[i][0];

        // Check chains and their leaders behave in the same way
        for (int i : c) {
            for (int j : c) {
                if (p[i][j] != p[chainLeader[myChain[i]]][j]) return 0;
            }
        }

        // There will be two ways to travel between a pair of nodes in different chains
        if (m == 2) return 0; // cannot have two chains (cannot make cycle)

        // Connect within chains
        for (int i : c) {
            if (chainLeader[myChain[i]] != i) {
                edge[i][chainLeader[myChain[i]]] = edge[chainLeader[myChain[i]]][i] = 1;
            }
        }

        // Connect chain leaders in a cycle
        if (m == 1) continue; // only one chain - no need
        for (int i = 0; i < m; ++i) {
            int j = (i+1)%m;
            edge[chainLeader[i]][chainLeader[j]] = edge[chainLeader[j]][chainLeader[i]] = 1;
        }
    }

    build(edge);
    return 1;
}