# Finding the nth ugly number

Posted on

Problem

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

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.

My working, accepted solution:

``````import heapq
class Solution:
def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
heap = 
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:
heapq.heappush(heap, num*p)

return heapq.heappop(heap)
``````

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

Solution

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 = 
seen = set()
while True:
num = heapq.heappop(heap)
yield num
for p in primes:
if num * p not in seen:
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))
``````

## 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 `return`s. `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:
``````to_add = {num*p for p in primes} - seen