Finding the nth ugly number 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]) -&gt; int:
        heap = [1]
        seen = set()
        while n&gt;=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?
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]) -&gt; int:
    # ...


class Solution:
    def nthSuperUglyNumber(self, n: int, primes: List[int]) -&gt; 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]) -&gt; 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]) -&gt; 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]) -&gt; int:
    ugly_numbers = super_ugly_numbers(primes)
    return next(itertools.islice(ugly_numbers, n - 1, n))