Find running median from a stream of integers

Posted on | In |

依序輸入多個數字，如何有效率找出中位數(median)？

例如：依序輸入7、13、2、6、9，中位數為7。

解法：

  1. 將所有數字排序，即可以得出中位數。
    時間複雜度為$O(nlogn)$

  2. 利用max heap和min heap：

    • 新增數字到heap：
      • 若max heap為空，數字新增到min heap。
      • 若max heap不為空並且數字小於max heap的root，數字新增到max heap。
      • 若max heap不為空並且數字大於max heap的root，數字新增到min heap。
    • rebalance：
      • 若max heap size - min heap size > 1, 刪除max heap的root並且新增到min heap。
      • 若min heap size - max heap size > 1, 刪除min heap的root並且新增到max heap。
    • 計算中位數：
      • 若max heap size > min heap size，中位數為max heap的root。
      • 若max heap size < min heap size，中位數為min heap的root。
      • 若max heap size == min heap size，中位數為$(max\ heap的root + min\ heap的root) / 2$ (根據定義決定)。

    時間複雜度為$O(1)$

    此方法是利用中位數的特性：

    • 左半部的數字一定小部右半邊。

      例如：2、3、7、13、15 或是2、3、7、13、15

    • 左半部和右半部的大小差值不能超過1，否則不能透過上述方法找出中位數。

    • 左半部的數字(包含中位數)建成max heap，max heap的root即為中位數。

      例如：2、3、7、13、15

      2、3、7建成max heap，13、15建成min heap

    • 右半部的數字(包含中位數)建成min heap，min heap的root即為中位數。

      例如：2、3、7、13、15

      2、3建成max heap，7、13、15建成min heap

程式碼：

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
62
63
64
65
66
67
68
69
70
71
72
73
74
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.PriorityQueue;
import java.util.Queue;

public class FindMedian<E> {
	private Queue<E>		queue;
	private Queue<E>		maxHeap;
	private Queue<E>		minHeap;
	private Comparator<E>	comparator;

	public FindMedian(Comparator<E> comparator) {
		queue = new LinkedList<>();
		maxHeap = new PriorityQueue<>(Collections.reverseOrder(comparator));
		minHeap = new PriorityQueue<>(comparator);
		this.comparator = comparator;
	}

	public void enqueue(E e) {
		if (e == null)
			throw new IllegalArgumentException();

		this.queue.add(e);

		addElementToHeap(e);

		rebalance();
	}

	public E dequeue() {
		E removeObject = this.queue.remove();
		this.maxHeap.remove(removeObject);
		this.minHeap.remove(removeObject);

		return removeObject;
	}

	public E peekMedian() {
		if (maxHeap.size() == minHeap.size()) {
			// Depends on definition
			return maxHeap.peek();
		}
		else if (maxHeap.size() > minHeap.size()) {
			return maxHeap.peek();
		}
		else {
			return minHeap.peek();
		}
	}

	private void addElementToHeap(E e) {
		E tmp = maxHeap.peek();
		if (tmp == null) {
			minHeap.add(e);
		}
		else if (comparator.compare(tmp, e) > 0) {
			maxHeap.add(e);
		}
		else {
			minHeap.add(e);
		}
	}

	private void rebalance() {
		int sizeDiff = maxHeap.size() - minHeap.size();
		if (sizeDiff > 1) {
			minHeap.add(maxHeap.remove());
		}
		else if (sizeDiff < -1) {
			maxHeap.add(minHeap.remove());
		}
	}
}

FindMedian

參考資料：

Find running median from a stream of integers