Problem: Redundant Connection II

In this problem, a rooted tree is a directed graph such that, there is exactly
one node (the root) for which all other nodes are descendants of this node,
plus every node has exactly one parent, except for the root node which has no
parents.

The given input is a directed graph that started as a rooted tree with `n`
nodes (with distinct values from `1` to `n`), with one additional directed
edge added. The added edge has two different vertices chosen from `1` to `n`,
and was not an edge that already existed.

The resulting graph is given as a 2D-array of `edges`. Each element of `edges`
is a pair `[ui, vi]` that represents a directed edge connecting nodes `ui` and
`vi`, where `ui` is a parent of child `vi`.

Return an edge that can be removed so that the resulting graph is a rooted
tree of `n` nodes. If there are multiple answers, return the answer that
occurs last in the given 2D-array.


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.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
    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 findRedundantDirectedConnection(self, edges: List[List[int]]) -> List[int]:
    ids = [0] * (len(edges) + 1)
    nodeWithTwoParents = 0

    for _, v in edges:
      ids[v] += 1
      if ids[v] == 2:
        nodeWithTwoParents = v

    def findRedundantDirectedConnection(skippedEdgeIndex: int) -> List[int]:
      uf = UnionFind(len(edges) + 1)

      for i, edge in enumerate(edges):
        if i == skippedEdgeIndex:
          continue
        if not uf.unionByRank(edge[0], edge[1]):
          return edge

      return []

    if nodeWithTwoParents == 0:
      return findRedundantDirectedConnection(-1)

    for i in reversed(range(len(edges))):
      _, v = edges[i]
      if v == nodeWithTwoParents:
        if not findRedundantDirectedConnection(i):
          return edges[i]


Java solution:
class UnionFind {
  public UnionFind(int 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];
    }
    return true;
  }

  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[] findRedundantDirectedConnection(int[][] edges) {
    int[] ids = new int[edges.length + 1];
    int nodeWithTwoParents = 0;

    for (int[] edge : edges) {
      final int v = edge[1];
      if (++ids[v] == 2) {
        nodeWithTwoParents = v;
        break;
      }
    }

    // If there is no edge with two ids, don't skip any edge.
    if (nodeWithTwoParents == 0)
      return findRedundantDirectedConnection(edges, -1);

    for (int i = edges.length - 1; i >= 0; --i)
      if (edges[i][1] == nodeWithTwoParents)
        // Try to delete the edges[i].
        if (findRedundantDirectedConnection(edges, i).length == 0)
          return edges[i];

    throw new IllegalArgumentException();
  }

  private int[] findRedundantDirectedConnection(int[][] edges, int skippedEdgeIndex) {
    UnionFind uf = new UnionFind(edges.length + 1);

    for (int i = 0; i < edges.length; ++i) {
      if (i == skippedEdgeIndex)
        continue;
      if (!uf.unionByRank(edges[i][0], edges[i][1]))
        return edges[i];
    }

    return new int[] {};
  }
}


C++ solution:
class UnionFind {
 public:
  UnionFind(int 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];
    }
    return true;
  }

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

class Solution {
 public:
  vector<int> findRedundantDirectedConnection(vector<vector<int>>& edges) {
    vector<int> ids(edges.size() + 1);
    int nodeWithTwoParents = 0;

    for (const vector<int>& edge : edges) {
      const int v = edge[1];
      if (++ids[v] == 2) {
        nodeWithTwoParents = v;
        break;
      }
    }

    // If there is no edge with two ids, don't skip any edge.
    if (nodeWithTwoParents == 0)
      return findRedundantDirectedConnection(edges, -1);

    for (int i = edges.size() - 1; i >= 0; --i)
      if (edges[i][1] == nodeWithTwoParents)
        // Try to delete the edges[i].
        if (findRedundantDirectedConnection(edges, i).empty())
          return edges[i];

    throw;
  }

  vector<int> findRedundantDirectedConnection(const vector<vector<int>>& edges,
                                              int skippedEdgeIndex) {
    UnionFind uf(edges.size() + 1);

    for (int i = 0; i < edges.size(); ++i) {
      if (i == skippedEdgeIndex)
        continue;
      if (!uf.unionByRank(edges[i][0], edges[i][1]))
        return edges[i];
    }

    return {};
  }
};