313. Super Ugly Number

package LeetCode_313

import java.util.*
import kotlin.collections.HashSet

/**
 * 313. Super Ugly Number
 * https://leetcode.com/problems/super-ugly-number/
 * A super ugly number is a positive integer whose prime factors are in the array primes.
Given an integer n and an array of integers primes, return the nth super ugly number.
The nth super ugly number is guaranteed to fit in a 32-bit signed integer.

Example 1:
Input: n = 12, primes = [2,7,13,19]
Output: 32
Explanation: [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].

Example 2:
Input: n = 1, primes = [2,3,5]
Output: 1
Explanation: 1 has no prime factors, therefore all of its prime factors are in the array primes = [2,3,5].

Constraints:
1. 1 <= n <= 106
2. 1 <= primes.length <= 100
3. 2 <= primes[i] <= 1000
4. primes[i] is guaranteed to be a prime number.
5. All the values of primes are unique and sorted in ascending order.
 * */
class Solution {
    /*
    * solution: use priority queue to keep the nth ugly number, and use Long to Maintain the accuracy of new ugly number;
    * Time complexity: O(n log(n*k)), k is the size of primes;
    * Space complexity: O(n)
    * */
    fun nthSuperUglyNumber(n: Int, primes: IntArray): Int {
        if (n == 1) {
            return 1
        }
        val queue = PriorityQueue<Long>()
        val set = HashSet<Long>()
        var result = 1L
        queue.offer(result)
        set.add(result)
        for (prime in primes) {
            queue.offer(prime.toLong())
            set.add(prime.toLong())
        }
        for (i in 0 until n) {//log(n)
            result = queue.poll()//log(k)
            for (prime in primes) {
                val multi:Long = prime.toLong() * result
                //contains of set run in O(1) time, contains of queue run in O(n) time
                if (!set.contains(multi)) {
                    set.add(multi)
                    queue.offer(multi)//log(k)
                }
            }
        }
        return result.toInt()
    }
}