forked from shuboc/LeetCode-2
-
Notifications
You must be signed in to change notification settings - Fork 1
/
super-ugly-number.py
142 lines (121 loc) · 4.05 KB
/
super-ugly-number.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
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
# Time: O(n * k)
# Space: O(n + k)
# Write a program to find the nth super ugly number.
#
# Super ugly numbers are positive numbers whose all
# prime factors are in the given prime list primes of size k.
# For example, [1, 2, 4, 7, 8, 13, 14, 16, 19, 26, 28, 32]
# is the sequence of the first 12 super ugly numbers given
# primes = [2, 7, 13, 19] of size 4.
#
# Note:
# (1) 1 is a super ugly number for any given primes.
# (2) The given numbers in primes are in ascending order.
# (3) 0 < k <= 100, 0 < n <= 106, 0 < primes[i] < 1000.
import heapq
# Heap solution. (620ms)
class Solution(object):
def nthSuperUglyNumber(self, n, primes):
"""
:type n: int
:type primes: List[int]
:rtype: int
"""
heap, uglies, idx, ugly_by_last_prime = [], [0] * n, [0] * len(primes), [0] * n
uglies[0] = 1
for k, p in enumerate(primes):
heapq.heappush(heap, (p, k))
for i in xrange(1, n):
uglies[i], k = heapq.heappop(heap)
ugly_by_last_prime[i] = k
idx[k] += 1
while ugly_by_last_prime[idx[k]] > k:
idx[k] += 1
heapq.heappush(heap, (primes[k] * uglies[idx[k]], k))
return uglies[-1]
# Time: O(n * k)
# Space: O(n + k)
# Hash solution. (932ms)
class Solution2(object):
def nthSuperUglyNumber(self, n, primes):
"""
:type n: int
:type primes: List[int]
:rtype: int
"""
uglies, idx, heap, ugly_set = [0] * n, [0] * len(primes), [], set([1])
uglies[0] = 1
for k, p in enumerate(primes):
heapq.heappush(heap, (p, k))
ugly_set.add(p)
for i in xrange(1, n):
uglies[i], k = heapq.heappop(heap)
while (primes[k] * uglies[idx[k]]) in ugly_set:
idx[k] += 1
heapq.heappush(heap, (primes[k] * uglies[idx[k]], k))
ugly_set.add(primes[k] * uglies[idx[k]])
return uglies[-1]
# Time: O(n * logk) ~ O(n * klogk)
# Space: O(n + k)
class Solution3(object):
def nthSuperUglyNumber(self, n, primes):
"""
:type n: int
:type primes: List[int]
:rtype: int
"""
uglies, idx, heap = [1], [0] * len(primes), []
for k, p in enumerate(primes):
heapq.heappush(heap, (p, k))
for i in xrange(1, n):
min_val, k = heap[0]
uglies += [min_val]
while heap[0][0] == min_val: # worst time: O(klogk)
min_val, k = heapq.heappop(heap)
idx[k] += 1
heapq.heappush(heap, (primes[k] * uglies[idx[k]], k))
return uglies[-1]
# Time: O(n * k)
# Space: O(n + k)
# TLE due to the last test case, but it passess and performs the best in C++.
class Solution4(object):
def nthSuperUglyNumber(self, n, primes):
"""
:type n: int
:type primes: List[int]
:rtype: int
"""
uglies = [0] * n
uglies[0] = 1
ugly_by_prime = list(primes)
idx = [0] * len(primes)
for i in xrange(1, n):
uglies[i] = min(ugly_by_prime)
for k in xrange(len(primes)):
if uglies[i] == ugly_by_prime[k]:
idx[k] += 1
ugly_by_prime[k] = primes[k] * uglies[idx[k]]
return uglies[-1]
# Time: O(n * logk) ~ O(n * klogk)
# Space: O(k^2)
# TLE due to the last test case, but it passess and performs well in C++.
class Solution5(object):
def nthSuperUglyNumber(self, n, primes):
"""
:type n: int
:type primes: List[int]
:rtype: int
"""
ugly_number = 0
heap = []
heapq.heappush(heap, 1)
for p in primes:
heapq.heappush(heap, p)
for _ in xrange(n):
ugly_number = heapq.heappop(heap)
for i in xrange(len(primes)):
if ugly_number % primes[i] == 0:
for j in xrange(i + 1):
heapq.heappush(heap, ugly_number * primes[j])
break
return ugly_number