long solve() {
long cost = 0;
Collections.sort(E);
// easy picks first
for (Edge e : E) {
if (K[e.a] && K[e.b]) {
cost += e.w;
e.w = 0;
}
}
// do the hard job
for (Edge e : E) {
if (e.w != 0) {
int temp = e.w;
e.w = 0;
int ma = countMachines(e.a);
int mb = countMachines(e.b);
if (ma > 0 && mb > 0) {
cost += temp;
}
}
}
return cost;
}
int countMachines(int start) {
int count = 0;
List<Edge> E2 = new ArrayList<Edge>();
Deque<Integer> D = new ArrayDeque<>();
boolean[] V = new boolean[G.size()];
V[start] = true;
D.add(start);
if (K[start]) {
count++;
}
while (!D.isEmpty()) {
int a = D.remove();
for (Edge e : G.get(a)) {
if (e.w == 0) {
continue;
}
;
E2.add(e);
int b = (e.a == a) ? e.b : e.a;
if (!V[b]) {
V[b] = true;
D.add(b);
if (K[b]) {
count++;
if (count > 1) {
return count;
}
}
}
}
}
if (count <= 1) {
for (Edge e : E2) {
e.w = 0;
}
}
return count;
}