LeetCode Solutions

632. Smallest Range Covering Elements from K Lists

Time: $O(n^2\log k)$, where $n = |\texttt{nums}|$

Space: $O(k)$

			

struct T {
  int i;
  int j;
  int num;  // nums[i][j]
  T(int i, int j, int num) : i(i), j(j), num(num) {}
};

class Solution {
 public:
  vector<int> smallestRange(vector<vector<int>>& nums) {
    auto compare = [&](const T& a, const T& b) { return a.num > b.num; };
    priority_queue<T, vector<T>, decltype(compare)> minHeap(compare);
    int mini = INT_MAX;
    int maxi = INT_MIN;

    for (int i = 0; i < nums.size(); ++i) {
      const int num = nums[i][0];
      minHeap.emplace(i, 0, num);
      mini = min(mini, num);
      maxi = max(maxi, num);
    }

    int minRange = mini;
    int maxRange = maxi;

    while (minHeap.size() == nums.size()) {
      const auto [i, j, _] = minHeap.top();
      minHeap.pop();
      if (j + 1 < nums[i].size()) {
        minHeap.emplace(i, j + 1, nums[i][j + 1]);
        maxi = max(maxi, nums[i][j + 1]);
        mini = minHeap.top().num;
        if (maxi - mini < maxRange - minRange) {
          minRange = mini;
          maxRange = maxi;
        }
      }
    }

    return {minRange, maxRange};
  }
};
			

class T {
  public int i;
  public int j;
  public int num; // nums[i][j]
  public T(int i, int j, int num) {
    this.i = i;
    this.j = j;
    this.num = num;
  }
}

class Solution {
  public int[] smallestRange(List<List<Integer>> nums) {
    Queue<T> minHeap = new PriorityQueue<>((a, b) -> a.num - b.num);
    int min = Integer.MAX_VALUE;
    int max = Integer.MIN_VALUE;

    for (int i = 0; i < nums.size(); ++i) {
      final int num = nums.get(i).get(0);
      minHeap.offer(new T(i, 0, num));
      min = Math.min(min, num);
      max = Math.max(max, num);
    }

    int minRange = min;
    int maxRange = max;

    while (minHeap.size() == nums.size()) {
      final int i = minHeap.peek().i;
      final int j = minHeap.poll().j;
      if (j + 1 < nums.get(i).size()) {
        minHeap.offer(new T(i, j + 1, nums.get(i).get(j + 1)));
        max = Math.max(max, nums.get(i).get(j + 1));
        min = minHeap.peek().num;
      }
      if (max - min < maxRange - minRange) {
        minRange = min;
        maxRange = max;
      }
    }

    return new int[] {minRange, maxRange};
  }
}
			

class Solution:
  def smallestRange(self, nums: List[List[int]]) -> List[int]:
    minHeap = [(row[0], i, 0) for i, row in enumerate(nums)]
    heapq.heapify(minHeap)

    maxRange = max(row[0] for row in nums)
    minRange = heapq.nsmallest(1, minHeap)[0][0]
    ans = [minRange, maxRange]

    while len(minHeap) == len(nums):
      num, r, c = heapq.heappop(minHeap)
      if c + 1 < len(nums[r]):
        heapq.heappush(minHeap, (nums[r][c + 1], r, c + 1))
        maxRange = max(maxRange, nums[r][c + 1])
        minRange = heapq.nsmallest(1, minHeap)[0][0]
        if maxRange - minRange < ans[1] - ans[0]:
          ans[0], ans[1] = minRange, maxRange

    return ans