使用优先队列寻找中位数

xuenima / 2023-05-16 / 原文

Next, Suppose we would like to invent a new ADT called MedianFinder which is a collection of integers and supports finding the median of the collection.

MedianFinder

add(x); // adds x to the collection of numbers

median(); // returns the median from a collection of numbers Describe

how you could implement this ADT by using existing Java ADTs as building blocks. What’s the most efficient implementation you can come up with?

import java.util.Comparator;
import java.util.PriorityQueue;

public class MedianFinder {
    private PriorityQueue<Integer> smallerHalf;
    private PriorityQueue<Integer> largerHalf;

    public MedianFinder() {
        smallerHalf = new PriorityQueue<>(Comparator.reverseOrder());
        largerHalf = new PriorityQueue<>();
    }

    public void add(int x) {
        if (smallerHalf.isEmpty() || x <= smallerHalf.peek()) {
            smallerHalf.add(x);
        } else {
            largerHalf.add(x);
        }
        balanceHeaps();
    }

    private void balanceHeaps() {
        if (smallerHalf.size() > largerHalf.size() + 1) {
            largerHalf.add(smallerHalf.poll());
        } else if (largerHalf.size() > smallerHalf.size()) {
            smallerHalf.add(largerHalf.poll());
        }
    }

    public double median() {
        if (smallerHalf.size() == largerHalf.size()) {
            return (smallerHalf.peek() + largerHalf.peek()) / 2.0;
        } else {
            return smallerHalf.peek();
        }
    }

    public static void main(String[] args) {
        MedianFinder test = new MedianFinder();
        test.add(8);
        test.add(7);
        test.add(6);
        test.add(5);
        test.add(4);
        test.add(3);
        test.add(2);
        double median = test.median();
        System.out.println(median);
    }
}

这段代码实现了一个名为 `MedianFinder` 的类，用于找到一系列整数的中位数。在这个类中，我们使用了两个优先队列（`PriorityQueue`），`smallerHalf` 和 `largerHalf` 分别存储较小的一半和较大的一半数字。优先队列可以很方便地获取、添加和删除元素。

1. `public MedianFinder()`: 构造函数，初始化两个优先队列。`smallerHalf` 为大顶堆（通过使用 `Comparator.reverseOrder()` 反转默认的小顶堆），存储较小的一半数字；`largerHalf` 为小顶堆，存储较大的一半数字。

2. `public void add(int x)`: 向集合中添加一个整数 `x`。根据 `x` 和 `smallerHalf` 堆顶元素的大小关系，将 `x` 添加到适当的堆中。然后调用 `balanceHeaps()` 方法来平衡两个堆。

3. `private void balanceHeaps()`: 平衡两个堆的大小。确保两个堆的大小之差不超过 1，这样中位数可以在常数时间内找到。如果 `smallerHalf` 的大小比 `largerHalf` 大 2 个或以上，那么将 `smallerHalf` 的堆顶元素移动到 `largerHalf`。反之，如果 `largerHalf` 的大小比 `smallerHalf` 大，那么将 `largerHalf` 的堆顶元素移动到 `smallerHalf`。

4. `public double median()`: 计算并返回中位数。如果两个堆的大小相同，那么中位数是两个堆顶元素的平均值；如果 `smallerHalf` 的大小比 `largerHalf` 大 1，那么中位数就是 `smallerHalf` 的堆顶元素。

5. `public static void main(String[] args)`: 主方法，用于测试 `MedianFinder` 类的功能。向 `MedianFinder` 对象中添加一些整数，然后计算并打印中位数。

这个实现利用了数据结构优先队列（`PriorityQueue`），可以在对数时间内向集合中添加整数，并在常数时间内找到中位数。这种方法在处理大量数据时非常高效。

下面是一个示例，展示了在类的成员变量中创建比较器的方式：

import java.util.Comparator;
import java.util.PriorityQueue;

public class CustomPriorityQueueExample {
    private Comparator<String> lengthComparator; // 类的成员变量

    public CustomPriorityQueueExample() {
        lengthComparator = new LengthComparator(); // 在构造函数中创建比较器实例
    }

    public void processElements() {
        PriorityQueue<String> pq = new PriorityQueue<>(lengthComparator);

        pq.add("Java");
        pq.add("Python");
        pq.add("C");
        pq.add("JavaScript");

        while (!pq.isEmpty()) {
            String element = pq.poll();
            System.out.println(element);
        }
    }

    private class LengthComparator implements Comparator<String> {
        @Override
        public int compare(String s1, String s2) {
            return Integer.compare(s1.length(), s2.length());
        }
    }

    public static void main(String[] args) {
        CustomPriorityQueueExample example = new CustomPriorityQueueExample();
        example.processElements();
    }
}