Update largest-component-size-by-common-factor.py

Author: kamyu104

Python/largest-component-size-by-common-factor.py

@@ -29,7 +29,7 @@ def largestComponentSize(self, A):
         :type A: List[int]
         :rtype: int
         """
-        def primeFactors(i):  # prime factor decomposition
+        def prime_factors(i):  # prime factor decomposition
             result = []
             d = 2
             if i%d == 0:
@@ -43,14 +43,67 @@ def primeFactors(i):  # prime factor decomposition
                         i //= d
                     result.append(d)
                 d += 2
-            if i > 2:
+            if i != 1:
                 result.append(i)
             return result
 
         union_find = UnionFind(len(A))
         nodesWithCommonFactor = collections.defaultdict(int)
         for i in xrange(len(A)):
-            for factor in primeFactors(A[i]):
+            for factor in prime_factors(A[i]):
+                if factor not in nodesWithCommonFactor:
+                    nodesWithCommonFactor[factor] = i
+                union_find.union_set(nodesWithCommonFactor[factor], i)
+        return max(union_find.size)
+
+
+# Time:  O(f * n), f is the max number of unique prime factors
+# Space: O(p + n), p is the total number of unique primes
+import collections
+
+
+class UnionFind(object):
+    def __init__(self, n):
+        self.set = range(n)
+        self.size = [1]*n
+
+    def find_set(self, x):
+        if self.set[x] != x:
+            self.set[x] = self.find_set(self.set[x])  # path compression.
+        return self.set[x]
+
+    def union_set(self, x, y):
+        x_root, y_root = map(self.find_set, (x, y))
+        if x_root == y_root:
+            return False
+        self.set[min(x_root, y_root)] = max(x_root, y_root)
+        self.size[max(x_root, y_root)] += self.size[min(x_root, y_root)]
+        return True
+
+
+class Solution2(object):
+    def largestComponentSize(self, A):
+        """
+        :type A: List[int]
+        :rtype: int
+        """
+        def prime_factors(x):  # prime factor decomposition
+            result = []
+            p = 2
+            while p*p <= x:
+                if x%p == 0:
+                    while x%p == 0:
+                        x //= p
+                    result.append(p)
+                p += 1
+            if x != 1:
+                result.append(x)
+            return result
+        
+        union_find = UnionFind(len(A))
+        nodesWithCommonFactor = collections.defaultdict(int)
+        for i in xrange(len(A)):
+            for factor in prime_factors(A[i]):
                 if factor not in nodesWithCommonFactor:
                     nodesWithCommonFactor[factor] = i
                 union_find.union_set(nodesWithCommonFactor[factor], i)