Median of K sorted array

Problem

Given k sorted arrays, find the median of the combined sorted array (without merging all arrays into one). The arrays may have different lengths.

Examples

Example 1

Input: arrays = [[1, 3, 5], [2, 4, 6]]
Output: 3.5
Explanation: The combined sorted array is [1, 2, 3, 4, 5, 6]. The median is (3 + 4) / 2 = 3.5.

Example 2

Input: arrays = [[1, 5, 9], [2, 3, 7, 10], [4, 6, 8]]
Output: 5.5
Explanation: The combined sorted array is [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]. The median is (5 + 6) / 2 = 5.5.

Solution

Method 1 – Min Heap (K-way Merge)

Intuition

Use a min heap to efficiently merge k sorted arrays and find the median without building the full merged array.

Approach

• Use a min heap to keep track of the smallest elements from each array.
• Extract the smallest element and push the next element from the same array into the heap.
• Repeat until reaching the median position(s).
• If the total number of elements is odd, the median is the middle element; if even, it is the average of the two middle elements.

Code

C++

class Solution {
public:
    double findMedian(vector<vector<int>>& arrays) {
        int n = 0;
        for (auto& arr : arrays) n += arr.size();
        auto cmp = [&](pair<int, pair<int, int>> a, pair<int, pair<int, int>> b) {
            return a.first > b.first;
        };
        priority_queue<pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, decltype(cmp)> minHeap(cmp);
        for (int i = 0; i < arrays.size(); ++i) {
            if (!arrays[i].empty()) minHeap.push({arrays[i][0], {i, 0}});
        }
        int count = 0, m1 = 0, m2 = 0;
        int mid1 = (n - 1) / 2, mid2 = n / 2;
        while (!minHeap.empty()) {
            auto [val, idx] = minHeap.top(); minHeap.pop();
            if (count == mid1) m1 = val;
            if (count == mid2) m2 = val;
            if (++count > mid2) break;
            int arrIdx = idx.first, elemIdx = idx.second + 1;
            if (elemIdx < arrays[arrIdx].size()) {
                minHeap.push({arrays[arrIdx][elemIdx], {arrIdx, elemIdx}});
            }
        }
        return (n % 2 == 0) ? (m1 + m2) / 2.0 : m2;
    }
};

Go

import "container/heap"

type Item struct {
    val, arrIdx, elemIdx int
}
type MinHeap []Item
func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i].val < h[j].val }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) { *h = append(*h, x.(Item)) }
func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

func FindMedian(arrays [][]int) float64 {
    n := 0
    for _, arr := range arrays { n += len(arr) }
    h := &MinHeap{}
    heap.Init(h)
    for i, arr := range arrays {
        if len(arr) > 0 {
            heap.Push(h, Item{arr[0], i, 0})
        }
    }
    count, m1, m2 := 0, 0, 0
    mid1, mid2 := (n-1)/2, n/2
    for h.Len() > 0 {
        item := heap.Pop(h).(Item)
        if count == mid1 { m1 = item.val }
        if count == mid2 { m2 = item.val }
        if count++; count > mid2 { break }
        if item.elemIdx+1 < len(arrays[item.arrIdx]) {
            heap.Push(h, Item{arrays[item.arrIdx][item.elemIdx+1], item.arrIdx, item.elemIdx+1})
        }
    }
    if n%2 == 0 {
        return float64(m1+m2) / 2.0
    }
    return float64(m2)
}

Java

class Solution {
    public double findMedian(List<List<Integer>> arrays) {
        int n = 0;
        for (List<Integer> arr : arrays) n += arr.size();
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        for (int i = 0; i < arrays.size(); i++) {
            if (!arrays.get(i).isEmpty()) minHeap.offer(new int[]{arrays.get(i).get(0), i, 0});
        }
        int count = 0, m1 = 0, m2 = 0;
        int mid1 = (n - 1) / 2, mid2 = n / 2;
        while (!minHeap.isEmpty()) {
            int[] top = minHeap.poll();
            if (count == mid1) m1 = top[0];
            if (count == mid2) m2 = top[0];
            if (++count > mid2) break;
            int arrIdx = top[1], elemIdx = top[2] + 1;
            if (elemIdx < arrays.get(arrIdx).size()) {
                minHeap.offer(new int[]{arrays.get(arrIdx).get(elemIdx), arrIdx, elemIdx});
            }
        }
        return (n % 2 == 0) ? (m1 + m2) / 2.0 : m2;
    }
}

Python

import heapq
class Solution:
    def find_median(self, arrays: list[list[int]]) -> float:
        n = sum(len(arr) for arr in arrays)
        min_heap = []
        for i, arr in enumerate(arrays):
            if arr:
                heapq.heappush(min_heap, (arr[0], i, 0))
        count, m1, m2 = 0, 0, 0
        mid1, mid2 = (n - 1) // 2, n // 2
        while min_heap:
            val, arr_idx, elem_idx = heapq.heappop(min_heap)
            if count == mid1: m1 = val
            if count == mid2: m2 = val
            if count > mid2: break
            count += 1
            if elem_idx + 1 < len(arrays[arr_idx]):
                heapq.heappush(min_heap, (arrays[arr_idx][elem_idx + 1], arr_idx, elem_idx + 1))
        return (m1 + m2) / 2 if n % 2 == 0 else m2

Complexity

• ⏰ Time complexity: O(N log k), where N is the total number of elements and k is the number of arrays (each heap operation is O(log k)).
• 🧺 Space complexity: O(k), for the heap.