Problem: Remove Max Number of Edges to Keep Graph Fully Traversable

Alice and Bob have an undirected graph of `n` nodes and three types of edges:

* Type 1: Can be traversed by Alice only.
* Type 2: Can be traversed by Bob only.
* Type 3: Can be traversed by both Alice and Bob.

Given an array `edges` where `edges[i] = [typei, ui, vi]` represents a
bidirectional edge of type `typei` between nodes `ui` and `vi`, find the
maximum number of edges you can remove so that after removing the edges, the
graph can still be fully traversed by both Alice and Bob. The graph is fully
traversed by Alice and Bob if starting from any node, they can reach all other
nodes.

Return the maximum number of edges you can remove, or return `-1` if Alice and
Bob cannot fully traverse the graph.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator

class UnionFind:
  def __init__(self, n: int):
    self.count = n
    self.id = list(range(n))
    self.rank = [0] * n

  def unionByRank(self, u: int, v: int) -> bool:
    i = self._find(u)
    j = self._find(v)
    if i == j:
      return False
    if self.rank[i] < self.rank[j]:
      self.id[i] = j
    elif self.rank[i] > self.rank[j]:
      self.id[j] = i
    else:
      self.id[i] = j
      self.rank[j] += 1
    self.count -= 1
    return True

  def _find(self, u: int) -> int:
    if self.id[u] != u:
      self.id[u] = self._find(self.id[u])
    return self.id[u]


class Solution:
  def maxNumEdgesToRemove(self, n: int, edges: List[List[int]]) -> int:
    alice = UnionFind(n)
    bob = UnionFind(n)
    requiredEdges = 0

    for type, u, v in sorted(edges, reverse=True):
      u -= 1
      v -= 1
      if type == 3:
        if alice.unionByRank(u, v) | bob.unionByRank(u, v):
          requiredEdges += 1
      elif type == 2:
        if bob.unionByRank(u, v):
          requiredEdges += 1
      else:
        if alice.unionByRank(u, v):
          requiredEdges += 1

    if alice.count == 1 and bob.count == 1:
        return len(edges) - requiredEdges
    else:
        return -1


Java solution:
class UnionFind {
  public UnionFind(int n) {
    count = n;
    id = new int[n];
    rank = new int[n];
    for (int i = 0; i < n; ++i)
      id[i] = i;
  }

  public boolean unionByRank(int u, int v) {
    final int i = find(u);
    final int j = find(v);
    if (i == j)
      return false;
    if (rank[i] < rank[j]) {
      id[i] = j;
    } else if (rank[i] > rank[j]) {
      id[j] = i;
    } else {
      id[i] = j;
      ++rank[j];
    }
    --count;
    return true;
  }

  public int getCount() {
    return count;
  }

  private int count;
  private int[] id;
  private int[] rank;

  private int find(int u) {
    return id[u] == u ? u : (id[u] = find(id[u]));
  }
}

class Solution {
  public int maxNumEdgesToRemove(int n, int[][] edges) {
    UnionFind alice = new UnionFind(n);
    UnionFind bob = new UnionFind(n);
    int requiredEdges = 0;

    // Greedily put type 3 edges in the front.
    Arrays.sort(edges, Comparator.comparingInt(edge -> - edge[0]));

    for (int[] edge : edges) {
      final int type = edge[0];
      final int u = edge[1] - 1;
      final int v = edge[2] - 1;
      switch (type) {
        case 3: // Can be traversed by Alice and Bob.
          // Note that we should use | instead of || because if the first
          // expression is true, short-circuiting will skip the second
          // expression.
          if (alice.unionByRank(u, v) | bob.unionByRank(u, v))
            ++requiredEdges;
          break;
        case 2: // Can be traversed by Bob.
          if (bob.unionByRank(u, v))
            ++requiredEdges;
          break;
        case 1: // Can be traversed by Alice.
          if (alice.unionByRank(u, v))
            ++requiredEdges;
      }
    }

    return alice.getCount() == 1 && bob.getCount() == 1 ? edges.length - requiredEdges : -1;
  }
}


C++ solution:
class UnionFind {
 public:
  UnionFind(int n) : count(n), id(n), rank(n) {
    iota(id.begin(), id.end(), 0);
  }

  bool unionByRank(int u, int v) {
    const int i = find(u);
    const int j = find(v);
    if (i == j)
      return false;
    if (rank[i] < rank[j]) {
      id[i] = j;
    } else if (rank[i] > rank[j]) {
      id[j] = i;
    } else {
      id[i] = j;
      ++rank[j];
    }
    --count;
    return true;
  }

  int getCount() const {
    return count;
  }

 private:
  int count;
  vector<int> id;
  vector<int> rank;

  int find(int u) {
    return id[u] == u ? u : id[u] = find(id[u]);
  }
};

class Solution {
 public:
  int maxNumEdgesToRemove(int n, vector<vector<int>>& edges) {
    UnionFind alice(n);
    UnionFind bob(n);
    int requiredEdges = 0;

    // Greedily put type 3 edges in the front.
    ranges::sort(edges, ranges::greater{},
                 [](const vector<int>& edge) { return edge[0]; });

    for (const vector<int>& edge : edges) {
      const int type = edge[0];
      const int u = edge[1] - 1;
      const int v = edge[2] - 1;
      switch (type) {
        case 3:  // Can be traversed by Alice and Bob.
          // Note that we should use | instead of || because if the first
          // expression is true, short-circuiting will skip the second
          // expression.
          if (alice.unionByRank(u, v) | bob.unionByRank(u, v))
            ++requiredEdges;
          break;
        case 2:  // Can be traversed by Bob.
          if (bob.unionByRank(u, v))
            ++requiredEdges;
          break;
        case 1:  // Can be traversed by Alice.
          if (alice.unionByRank(u, v))
            ++requiredEdges;
          break;
      }
    }

    return alice.getCount() == 1 && bob.getCount() == 1
               ? edges.size() - requiredEdges
               : -1;
  }
};