leetcode-union-find

128. Longest Consecutive Sequence.py

class UnionFind:
    def __init__(self, n):
        self._parents = [node for node in range(n)]
        self._rank = [1 for _ in range(n)]
        
    def find(self, u):
        while u != self._parents[u]:
            self._parents[u] = self._parents[self._parents[u]]
            u = self._parents[u]
        return u 
    
    def union(self, u, v):
        root_u, root_v = self.find(u), self.find(v)
        if root_u == root_v:
            return True
        if self._rank[root_u] > self._rank[root_v]:
            self._parents[root_v] = root_u
        elif self._rank[root_u] < self._rank[root_v]:
            self._parents[root_u] = root_v
        else:
            self._parents[root_v] = root_u
            self._rank[root_u] +=1
        return False
    
class Solution:
    """
    PC:
        duplicate?
        
    [1,2,0,1]
    重複的元素不影響我們的答案可以去重
    [1,2,0]
    
    Ref:
        https://leetcode.com/problems/longest-consecutive-sequence/discuss/263634/Python-Union-Find
    """
    def longestConsecutive(self, nums: List[int]) -> int:
        """
        O(nlogn)
        """
        nums.sort()
        n = len(nums)
        if n == 0:
            return 0
        if n == 1:
            return 1
        l = 1
        max_len = -1
        for i in range(1, n):
            if nums[i] ==nums[i-1]+1:
                l += 1
            elif nums[i] ==nums[i-1]:
                l = l
            else:
                # reset length, 重新計算
                l = 1
            max_len = max(max_len, l)
        return max_len
    
    def longestConsecutive(self, nums: List[int]) -> int:
        nums = set(nums)
        n = len(nums)
        val2index = {}
        for i, num in enumerate(nums):
            val2index[num] = i
        
        union_find = UnionFind(n)
        for i, num in enumerate(nums):
            # 如果差距只差1的node, 我們就把他接起來, 每個node用index來代表
            if num - 1 in val2index:
                union_find.union(i, val2index[num-1])
            if num + 1 in val2index:
                union_find.union(i, val2index[num+1])
        
        count = collections.defaultdict(int)
        for i, num in enumerate(nums):
            root = union_find.find(i)
            count[root] += 1
        ans = 0
        for root, nb_connected in count.items():
            ans = max(ans, nb_connected)
        
        return ans