JAVA笔试面试之使用堆排对大量数据进行筛选最大或最小

JAVA笔试或面试时经常会碰到的一个题目就是用JAVA处理大量数据，筛选出最大或最小的N（相对的数量级较小，如100）个数据，要求尽可能考虑时间复杂度和空间复杂度。使用堆排是一种比较好的算法，另一种思路是利用快排时用到的partition函数（不过partition函数会对原有数据的顺序进行修改），下面是我写的用堆排筛选最大或最小的实现代码：

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

/**
 * Created by violetMoon on 2016/4/20.
 */
public class HeapSort <T extends Comparable<? super T>> {

    private List<T> sortResult;
    private int maxNum;
    private boolean isSelectMax;

    public HeapSort(int maxNum) {
        this(maxNum, true);
    }

    public HeapSort(int maxNum, boolean selectMax) {
        sortResult = new ArrayList<>(maxNum);
        this.maxNum = maxNum;
        isSelectMax = selectMax;
    }

    public void sort(T[] datas) {
        for (T data: datas)
            insert(data);
    }

    private void insert(T data) {
        if (sortResult.size() == 0) {
            sortResult.add(data);
            return;
        }

        T min = sortResult.get(0);
        if (sortResult.size() == maxNum) {
            int cmp = data.compareTo(min);
            if (isSelectMax) {
                if (cmp < 0) //选最大的时候需要维护一个最小堆
                    return;
                else {
                    sortResult.set(0, data);
                    pearDown(0);
                }
            } else {
                if (cmp > 0) //选最小的时候需要维护一个最大堆
                    return;
                else {
                    sortResult.set(0, data);
                    pearDown(0);
                }
            }
        } else {
            sortResult.add(data);
            pearUp(sortResult.size() - 1);
        }
    }

    private void pearUp(int index) {
        int curIndex = index;
        T upData = sortResult.get(curIndex);
        while (curIndex != 0) {
            int fatherIndex = (curIndex - 1) / 2;
            T father = sortResult.get(fatherIndex);
            int cmp = father.compareTo(upData);
            if (isSelectMax) {
                if (cmp > 0) { //选最大的时候需要维护一个最小堆
                    sortResult.set(curIndex, father);
                    curIndex = fatherIndex;
                    continue;
                } else
                    break;
            } else {
                if (cmp < 0) { //选最小的时候需要维护一个最大堆
                    sortResult.set(curIndex, father);
                    curIndex = fatherIndex;
                    continue;
                } else
                    break;
            }
        }
        sortResult.set(curIndex, upData);
    }

    private void pearDown(int index) {
        int curIndex = index;
        T downData = sortResult.get(curIndex);
        while (true) {
            int leftChildIndex = 2 * curIndex + 1;
            T leftChild = leftChildIndex < sortResult.size()? sortResult.get(leftChildIndex): null;
            int rightChildIndex = 2 * curIndex + 2;
            T rightChild = rightChildIndex < sortResult.size()? sortResult.get(rightChildIndex): null;
            int candidateChildIndex;
            T candidateChild;
            if (leftChild != null && rightChild != null) {
                if (isSelectMax) {
                    candidateChild = leftChild.compareTo(rightChild) < 0 ? leftChild : rightChild;
                    candidateChildIndex = candidateChild == leftChild ? leftChildIndex : rightChildIndex;
                } else {
                    candidateChild = leftChild.compareTo(rightChild) > 0 ? leftChild : rightChild;
                    candidateChildIndex = candidateChild == leftChild ? leftChildIndex : rightChildIndex;
                }
            } else if (leftChild != null && rightChild == null) {
                candidateChildIndex = leftChildIndex;
                candidateChild = leftChild;
            } else {
                break;
            }

            int cmp = downData.compareTo(candidateChild);
            if (isSelectMax) {
                if (cmp > 0) { //选最大的时候需要维护一个最小堆
                    sortResult.set(curIndex, candidateChild);
                    curIndex = candidateChildIndex;
                    continue;
                } else {
                    break;
                }
            } else {
                if (cmp < 0) { //选最小的时候需要维护一个最大堆
                    sortResult.set(curIndex, candidateChild);
                    curIndex = candidateChildIndex;
                    continue;
                } else {
                    break;
                }
            }
        }

        sortResult.set(curIndex, downData);
    }

    public boolean hasNext() {
        return sortResult.size() != 0;
    }

    public T getNext() {
        T min = sortResult.get(0);
        if (sortResult.size() > 1) {
            int size = sortResult.size();
            sortResult.set(0, sortResult.get(size - 1));
            sortResult.remove(size - 1);
            pearDown(0);
        } else
            sortResult.clear();
        return min;
    }

    public static void main(String[] args) {
        HeapSort<Integer> sorter = new HeapSort<>(100, false);
        Integer[] datas = new Integer[1000];

        Random random = new Random(System.currentTimeMillis());
        for (int i=0; i<datas.length; ++i)
            datas[i] = random.nextInt(100000);

        sorter.sort(datas);

        for (int i=0; i<datas.length; ++i)
            datas[i] = random.nextInt(100000);
        sorter.sort(datas);

        while (sorter.hasNext())
            System.out.println(sorter.getNext());
    }
}

main函数中只是简单地测试了下，实际使用中为了达到较好的IO和CPU效率可能需要开一些IO线程和一个排序用的CPU线程（因为HeapSort中只考虑在单线程环境下进行排序）