Update people-whose-list-of-favorite-companies-is-not-a-subset-of-another-list.py

Author: kamyu104

Python/people-whose-list-of-favorite-companies-is-not-a-subset-of-another-list.py

@@ -1,7 +1,7 @@
-# Time:  O(n * c * l + n^2 * c), n is favoriteCompanies.length
-#                              , c is the max of favoriteCompanies[i].length
+# Time:  O(n * m * l + n^2 * m), n is favoriteCompanies.length
+#                              , m is the max of favoriteCompanies[i].length
 #                              , l is the max of favoriteCompanies[i][j].length
-# Space: O(n * c * l)
+# Space: O(n * m * l)
 
 class Solution(object):
     def peopleIndexes(self, favoriteCompanies):
@@ -22,10 +22,10 @@ def peopleIndexes(self, favoriteCompanies):
 
 
 
-# Time:  O(n^2 * c * l * log*(n)), n is favoriteCompanies.length
-#                                , c is the max of favoriteCompanies[i].length
-#                                , l is the max of favoriteCompanies[i][j].length
-# Space: O(n * c * l)
+# Time:  O(n * m * l + n^2 * m * log*(n)), n is favoriteCompanies.length
+#                                        , m is the max of favoriteCompanies[i].length
+#                                        , l is the max of favoriteCompanies[i][j].length
+# Space: O(n * m * l)
 class UnionFind(object):
     def __init__(self, data):
         self.data = [set(d) for d in data]
@@ -54,9 +54,16 @@ def peopleIndexes(self, favoriteCompanies):
         :type favoriteCompanies: List[List[str]]
         :rtype: List[int]
         """
-        union_find = UnionFind(favoriteCompanies)
-        for i in xrange(len(favoriteCompanies)):
-            for j in xrange(len(favoriteCompanies)):
+        lookup, comps = {}, []
+        for cs in favoriteCompanies:
+            comps.append(set())
+            for c in cs:
+                if c not in lookup:
+                    lookup[c] = len(lookup)
+                comps[-1].add(lookup[c])
+        union_find = UnionFind(comps)
+        for i in xrange(len(comps)):
+            for j in xrange(len(comps)):
                 if j == i:
                     continue
                 union_find.union_set(i, j)