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.
这题算是Ugly Number II的一个拓展,主要是这里生成丑数的因数不再是[2,3,5]而是一个素数数组.其实做法还可以跟以前一样,只不过用一个和素数数组等长的数组来保存每个素数对应的index.产生一个素数时,更新这个数组.更新的判读方式也和之前一样,这样的复杂为O(nk).主要的消耗在于每次需要从k个素数生成的备选中,选择一个最小值,和更新index这个是O(k)的复杂度.那么有没有办法改进呢,有,可以看到每次前面的做法每次产生一个丑数,都要用每个素数产生一个备选,这些备选本身极其不可能一次全部更新,所以每次的候选中有相当大的重复部分.我们可以把备选放入heap当中.这样不用每次全部比较得到最小值(heap里面本身已经保存了相对关系).这样做的时间复杂度为O(nlogk),具体怎么证明还有待思考.附上一个很好的discuss
def nthSuperUglyNumber(self, n, primes): """ :type n: int :type primes: List[int] :rtype: int """ length = len(primes) uglys = [1] times = [0] * length #prime[i]'s multipy places in uglys minHeap = [(primes[i],i) for i in xrange(length)] heapq.heapify(minHeap) while len(uglys) < n: (val, index) = minHeap[0] uglys.append(val) while minHeap[0][0] == uglys[-1]: index = heapq.heappop(minHeap)[1] times[index] += 1 heapq.heappush(minHeap,(primes[index]*uglys[times[index]],index)) return uglys[-1]