def test_union_find(self): partners = [0, 1, 3, 2] _ = MegaStopFac(0) self.assertListEqual(_.union_find(partners), [0, 1, 2, 2]) partners = [0, 1, 2, 3] self.assertListEqual(_.union_find(partners), partners)
def test_get_groups(self): partners = [0, 1, 3, 2] _ = MegaStopFac(0) groups = _.union_find(partners) self.assertDictEqual(_.get_groups(groups), {0: [0], 1: [1], 2: [2, 3]})