LeetCode 264. Ugly Number II

Description

https://leetcode.com/problems/ugly-number-ii/

An ugly number is a positive integer whose prime factors are limited to 2, 3, and 5.

Given an integer n, return the nth ugly number.

Example 1:

Input: n = 10
Output: 12
Explanation: [1, 2, 3, 4, 5, 6, 8, 9, 10, 12] is the sequence of the first 10 ugly numbers.

Example 2:

Input: n = 1
Output: 1
Explanation: 1 has no prime factors, therefore all of its prime factors are limited to 2, 3, and 5.

Constraints:

• 1 <= n <= 1690

Python Solution

Use a heap to maintain the ugly numbers and to find the nth ugly number.

class Solution:
    def nthUglyNumber(self, n: int) -> int:
        heap = [1]

        visited = set([1])

        current_ugly = 1

        for _ in range(n):

            current_ugly = heapq.heappop(heap)

            for factor in [2, 3, 5]:
                new_ugly = current_ugly * factor

                if new_ugly not in visited:
                    visited.add(new_ugly)
                    heapq.heappush(heap, new_ugly)

        return current_ugly