forked from HuichuanLI/play-with-data-structure-python
-
Notifications
You must be signed in to change notification settings - Fork 1
/
UnionFind1.py
31 lines (25 loc) · 817 Bytes
/
UnionFind1.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from base import UF
# First version - quick find
class UnionFind1(UF):
def __init__(self, size):
self._id = [i for i in range(size)]
def get_size(self):
return len(self._id)
def _find(self, p):
if p < 0 or p >= len(self._id):
raise ValueError('p is out of bound.')
return self._id[p]
# O(1)
def is_connected(self, p, q):
return self._find(p) == self._find(q)
# O(n)
def union_elements(self, p, q):
if p < 0 or p >= len(self._id) or q < 0 or q >= len(self._id):
raise ValueError('Illegal argument.')
p_id = self._find(p)
q_id = self._find(q)
if p_id == q_id:
return
for i in range(len(self._id)):
if self._id[i] == p_id:
self._id[i] = q_id