Skip to content

313. Super Ugly Number 👍

Approach 1: Similar to 264. Ugly Number II

  • Time: $O(nk)$
  • Space: $O(k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
 public:
  int nthSuperUglyNumber(int n, vector<int>& primes) {
    const int k = primes.size();
    vector<int> indices(k);
    vector<int> uglyNums{1};

    while (uglyNums.size() < n) {
      vector<int> nexts(k);
      for (int i = 0; i < k; ++i)
        nexts[i] = uglyNums[indices[i]] * primes[i];
      const int next = ranges::min(nexts);
      for (int i = 0; i < k; ++i)
        if (next == nexts[i])
          ++indices[i];
      uglyNums.push_back(next);
    }

    return uglyNums.back();
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
  public int nthSuperUglyNumber(int n, int[] primes) {
    final int k = primes.length;
    int[] indices = new int[k];
    int[] uglyNums = new int[n];
    uglyNums[0] = 1;

    for (int i = 1; i < n; ++i) {
      int[] nexts = new int[k];
      for (int j = 0; j < k; ++j)
        nexts[j] = uglyNums[indices[j]] * primes[j];
      final int next = Arrays.stream(nexts).min().getAsInt();
      for (int j = 0; j < k; ++j)
        if (next == nexts[j])
          ++indices[j];
      uglyNums[i] = next;
    }

    return uglyNums[n - 1];
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution:
  def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
    k = len(primes)
    nums = [1]
    indices = [0] * k

    while len(nums) < n:
      nexts = [0] * k
      for i in range(k):
        nexts[i] = nums[indices[i]] * primes[i]
      next = min(nexts)
      for i in range(k):
        if next == nexts[i]:
          indices[i] += 1
      nums.append(next)

    return nums[-1]

Approach 2: Heap

  • Time: $O(n\log k)$
  • Space: $O(k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
struct UglyNum {
  int prime;
  int index;   // the next index in `uglyNums`
  long value;  // prime * uglyNums[index]
  UglyNum(int prime, int index, long value)
      : prime(prime), index(index), value(value) {}
};

class Solution {
 public:
  int nthSuperUglyNumber(int n, vector<int>& primes) {
    auto compare = [&](const UglyNum& a, const UglyNum& b) {
      return a.value > b.value;
    };
    priority_queue<UglyNum, vector<UglyNum>, decltype(compare)> minHeap(
        compare);
    vector<int> uglyNums{1};

    for (const int prime : primes)
      minHeap.emplace(prime, 1, prime * uglyNums[0]);

    while (uglyNums.size() < n) {
      uglyNums.push_back(minHeap.top().value);
      while (minHeap.top().value == uglyNums.back()) {
        const auto [prime, index, _] = minHeap.top();
        minHeap.pop();
        minHeap.emplace(prime, index + 1,
                        prime * static_cast<long>(uglyNums[index]));
      }
    }

    return uglyNums.back();
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class UglyNum {
  public int prime;
  public int index; // the next index in `uglyNums`
  public int value; // prime * uglyNums[index]
  public UglyNum(int prime, int index, int value) {
    this.prime = prime;
    this.index = index;
    this.value = value;
  }
}

class Solution {
  public int nthSuperUglyNumber(int n, int[] primes) {
    Queue<UglyNum> minHeap = new PriorityQueue<>((a, b) -> Integer.compare(a.value, b.value));
    int[] uglyNums = new int[n];
    uglyNums[0] = 1;

    for (final int prime : primes)
      minHeap.offer(new UglyNum(prime, 1, prime * uglyNums[0]));

    for (int i = 1; i < n; ++i) {
      uglyNums[i] = minHeap.peek().value;
      while (minHeap.peek().value == uglyNums[i]) {
        final UglyNum u = minHeap.poll();
        minHeap.offer(new UglyNum(u.prime, u.index + 1, u.prime * uglyNums[u.index]));
      }
    }

    return uglyNums[n - 1];
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class UglyNum:
  def __init__(self, prime: int, index: int, value: int):
    self.prime = prime
    self.index = index  # Point the next index of uglyNums.
    self.value = value  # prime * uglyNums[index]


class Solution:
  def nthSuperUglyNumber(self, n: int, primes: list[int]) -> int:
    minHeap = []  # (value, prime, index)
    uglyNums = [1]

    for prime in primes:
      heapq.heappush(minHeap, (prime * uglyNums[0], prime, 1))

    while len(uglyNums) < n:
      uglyNums.append(minHeap[0][0])
      while minHeap[0][0] == uglyNums[-1]:
        _, prime, index = heapq.heappop(minHeap)
        heapq.heappush(minHeap, (prime * uglyNums[index], prime, index + 1))

    return uglyNums[-1]