text: # Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved
  """Simple union find structure implementation"""
  
  
  class UnionFind:
      """Optimized union find structure"""
  
      def __init__(self, n):
          """Initialize a union find with n components"""
          self.compo = list(range(n))
          self.weight = [1] * n
          self.nb_compo = n
  
      def get_nb_compo(self):
          return self.nb_compo
  
      def find(self, x):
          if self.compo[x] == x:
              return x
          self.compo[x] = self.find(self.compo[x])
          return self.compo[x]
  
      def unite(self, a, b):
          fa = self.find(a)
          fb = self.find(b)
          if fa != fb:
              self.nb_compo -= 1
              if self.weight[fb] > self.weight[fa]:
                  fa, fb = fb, fa
              self.compo[fb] = fa
              self.weight[fa] += self.weight[fb]
  