PROBLEM DESCRIPTION
Given an input stream of N integers. The task is to insert these numbers into a new stream and find the median of the stream formed by each insertion of X to the new stream.
SOLUTION
We can solve this by using two heaps: a max-heap to store the smaller half of the numbers and a min-heap to store the larger half. Each time we insert a new number, we add it to the appropriate heap and balance the heaps to ensure their sizes stay nearly equal. The median is then calculated based on the sizes of the heaps: if both heaps have the same size, the median is the average of the top elements from both heaps; if one heap is larger, the median is the top element of that heap.
See this for detailed explanation:
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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class Solution {
// Two priority queues (heaps) to maintain the two halves of the numbers seen so far.
// maxHeap will store the smaller half of the numbers (maximum element at the top).
// minHeap will store the larger half of the numbers (minimum element at the top).
static Queue<Integer> maxHeap;
static Queue<Integer> minHeap;
Solution() {
maxHeap = new PriorityQueue<Integer>(Collections.reverseOrder());
minHeap = new PriorityQueue<Integer>();
}
// Function to insert a new number into the appropriate heap.
public static void insertHeap(int x) {
// If maxHeap is empty, directly add the number to maxHeap.
if (maxHeap.isEmpty()) {
maxHeap.add(x);
} else {
// If the number is greater than the top of maxHeap, it belongs in minHeap.
if (x > maxHeap.peek()) {
minHeap.add(x);
} else {
// Otherwise, add it to maxHeap.
maxHeap.add(x);
}
}
// Balance the heaps after the insertion.
balanceHeaps();
}
// Function to balance the two heaps.
public static void balanceHeaps() {
// If maxHeap has more than 1 extra element compared to minHeap, move the top element of maxHeap to minHeap.
if (maxHeap.size() - minHeap.size() > 1) {
minHeap.add(maxHeap.poll());
}
// If minHeap has more elements than maxHeap, move the top element of minHeap to maxHeap.
else if (minHeap.size() - maxHeap.size() > 0) {
maxHeap.add(minHeap.poll());
}
}
// Function to return the median of the current stream.
public static double getMedian() {
// Calculate the total number of elements in both heaps.
int n = maxHeap.size() + minHeap.size();
// If the total number of elements is even, the median is the average of the top elements from both heaps.
if (n % 2 == 0) {
return (maxHeap.peek() + minHeap.peek()) / 2.0;
} else {
// If the total number of elements is odd, the median is the top element of maxHeap.
return maxHeap.peek() / 1.0;
}
}
}