Finding the nth ugly number

Question

I'm trying to optimize my solution to this LeetCode problem:

Write a program to find the n^th super ugly number.

Super ugly numbers are positive numbers whose all prime factors are in the given prime list primes of size k.

My working, accepted solution:

import heapq
class Solution:
 def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
 heap = [1]
 seen = set()
 while n>=0:
 num = heapq.heappop(heap)
 n -=1
 if n == 0:
 return num 
 for p in primes:
 if num * p not in seen:
 seen.add(num*p)
 heapq.heappush(heap, num*p)
 return heapq.heappop(heap)

What are some ways to optimize it in terms of time and space complexity?

Answer 1

As already explained in this answer, there is no need to use a class with an instance method here. In addition, the convention for function names in Python is to use snake_case, not camelCase.

On the other hand, this template is given on the LeetCode submission template, so you are not free to change it.

What you can do is to make the required solution method a wrapper for a free function (which then is universally usable):

def nth_super_ugly_number(n: int, primes: List[int]) -> int:
 # ...
class Solution:
 def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
 return nth_super_ugly_number(n, primes)

Generating the ugly numbers can be separated from counting them (until the nth ugly number is reached) if the former is done with a generator. The generator is now a reusable function as well, and you might want to add a docstring comment:

def super_ugly_numbers(primes: List[int]) -> int:
 """Generate super ugly numbers.
 Super ugly numbers are positive integers whose all prime factors are in
 the given prime list.
 """
 heap = [1]
 seen = set()
 while True:
 num = heapq.heappop(heap)
 yield num
 for p in primes:
 if num * p not in seen:
 seen.add(num * p)
 heapq.heappush(heap, num * p)
def nth_super_ugly_number(n: int, primes: List[int]) -> int:
 ugly_numbers = super_ugly_numbers(primes)
 for _ in range(n - 1):
 next(ugly_numbers)
 return next(ugly_numbers)

The last function can be simplified using itertools.islice:

import itertools
def nth_super_ugly_number(n: int, primes: List[int]) -> int:
 ugly_numbers = super_ugly_numbers(primes)
 return next(itertools.islice(ugly_numbers, n - 1, n))

Answer 2

Loop like a native

 while n>=0:
 n -=1

You don't actually use n, nor should you be manually decrementing it, so this can become

for _ in range(n+1):

Returns

I'm not clear on why you have two returns. return num will take effect on the last iteration of the loop, so how would the other return happen at all? Probably you should remove the if n == 0 condition and rework the place where heappop is called so that the for _ in range can completely control the loop.

Set operations

Try this out to see what performance impact it has:

Rather than

 for p in primes:
 if num * p not in seen:
 seen.add(num*p)
 heapq.heappush(heap, num*p)

try

to_add = {num*p for p in primes} - seen
seen |= to_add
for a in to_add:
 heapq.heappush(heap, a)