#include "supertrees.h"
#include <bits/stdc++.h>
using namespace std;
using vi = vector<int>;
using vvi = vector<vi>;
struct DSU{
vi p, s;
DSU(int size){
p = s = vi(size + 1, 1);
for(int i = 0; i <= size; i++)
p[i] = i;
}
int find(int x){
if(p[x] == x)
return x;
return p[x] = find(p[x]);
}
bool same(int a, int b){
return find(a) == find(b);
}
void unite(int a, int b){
a = find(a);
b = find(b);
if(a == b)
return;
if(s[a] > s[b])
swap(a, b);
s[b] += s[a];
p[a] = b;
}
};
/*
0 [X, 1, 2, 2]
1 [1, X, 2, 2]
2 [2, 2, X, 2]
3 [2, 2, 2, X]
X 0 1 2 3
*/
int construct(vvi p) {
int n = p.size();
vvi answer = vvi(n, vi(n, 0));
//connect 1s
DSU dsu1(n);
set<int> red;
for (int i = 0; i < n; i++)
for (int j = i + 1; j < n; j++)
if(p[i][j] == 1){
dsu1.unite(i, j);
red.insert(i);
red.insert(j);
}
vvi group1(n, vi());
for (int i = 0; i < n; i++)
group1[dsu1.find(i)].push_back(i);
for(vi & g : group1){
if(g.size() <= 1)
continue;
for (int i = 0; i + 1 < (int)g.size(); i++)
{
int cur = g[i], nxt = g[i + 1];
answer[cur][nxt] = 1;
answer[nxt][cur] = 1;
}
}
//connect 2s
DSU dsu2(n);
set<int> orange;
for (int i = 0; i < n; i++)
for (int j = i + 1; j < n; j++){
if(red.find(i) != red.end())
continue;
if(red.find(j) != red.end())
continue;
if(p[i][j] == 2){
dsu2.unite(i, j);
orange.insert(i);
orange.insert(j);
}
}
vvi group2(n, vi());
for (int i = 0; i < n; i++)
group2[dsu2.find(i)].push_back(i);
for(vi & g : group2){
if(g.size() <= 1)
continue;
for (int i = 0; i + 1 < (int)g.size(); i++)
{
int cur = g[i], nxt = g[i + 1];
answer[cur][nxt] = 1;
answer[nxt][cur] = 1;
}
}
DSU dsuGroups(n);
for(int r : red)
for(int o : orange){
if(p[r][o] == 0)
continue;
int rGroupId = dsu1.find(r);
int oGroupId = dsu2.find(o);
if(dsuGroups.same(rGroupId, oGroupId))
continue;
dsuGroups.unite(rGroupId, oGroupId);
vi & redGroup = group1[rGroupId];
vi & orangeGroup = group2[oGroupId];
int redBack = redGroup.back();
int orangeFront = orangeGroup[0];
int orangeBack = orangeGroup.back();
answer[redBack][orangeBack] = 1;
answer[orangeBack][redBack] = 1;
answer[redBack][orangeFront] = 1;
answer[orangeFront][redBack] = 1;
}
// for(int x : red)
// cout << x << " ";
// cout << '\n';
// for(int x : orange)
// cout << x << " ";
// cout << '\n';
build(answer);
return 1;
}
