堆原理分析及优先级队列实现

前言

堆是一颗具有特定性质的二叉树，基于堆可以实现优先级队列，在搜索，插入和删除操作上的时间复杂度均为O(logn)，在找到最大或最小元素操作上的时间复杂度均为O(1)。本篇文章将对堆的原理和性质进行分析，并提供基于JAVA语言的堆的实现，最后会再基于堆，实现一个简单优先级队列。

正文

一. 堆的定义

堆的定义如下所示。

• 堆是一颗完全二叉树；
• 父节点的元素如果总是大于等于子节点，称之为大根堆；
• 父节点的元素如果总是小于等于子节点，称之为小根堆。

要理解堆，首先需要了解什么是完全二叉树。完全二叉树定义为：如果二叉树的深度为h，那么除第h层外，其余第1到h-1层的节点数都达到最大值，第h层的所有节点都连续集中在左边。一颗完全二叉树如下所示。

大根堆如下所示。

小根堆如下所示。

通常，堆中元素用一个数组来存放。以上图的小根堆为例，节点旁边的数字代表当前节点在数组中的索引，可以发现其规律就是：按照层序遍历的方式遍历完全二叉树，并每遍历一个节点就将当前节点的元素添加到数组中。所以上图的小根堆以数组表示如下。

最后，可以根据堆中某个节点在数组中的索引来计算这个节点的左右子节点和父节点在数组中的索引。计算的JAVA代码如下所示。

public class HeapTry {

    public int getLeft(int i) {
        return (i + 1) * 2 - 1;
    }

    public int getRigth(int i) {
        return (i + 1) * 2;
    }

    public int getParent(int i) {
        if (i == 0) {
            return 0;
        }
        return (i - 1) / 2;
    }
    
    // 省略其它

}

二. 堆的性质

• 大根堆的性质：父节点的元素总是大于等于子节点
• 小根堆的性质：父节点的元素总是小于等于子节点

以小根堆为例，考虑这样一个场景：给定一个小根堆，然后将根节点的元素删除并为根节点添加一个新的元素，新的元素不满足小于等于子节点的元素，此时小根堆的性质遭到破坏。图示如下。

对于上述场景，如果需要恢复小根堆的性质，则以根节点作为第一个父节点，按照如下步骤执行。

1. 将父节点与子节点的元素进行比较并得到元素最小的节点，如果元素最小的节点是父节点，那么已经恢复小根堆的性质，执行结束；如果元素最小的节点是某个子节点，执行步骤2；
2. 将父节点与元素最小的子节点互换元素，由于互换元素后子节点有可能破坏小根堆性质，因此从子节点开始，执行步骤1，直到恢复小根堆性质。

图示如下。

将恢复小根堆性质的算法用JAVA代码实现如下。

public class HeapTry {

    // 省略其它

    public void keepHeap(int[] array, int i, int size) {
        // 通过父节点计算左右子节点的数组索引
        int left = getLeft(i);
        int rigth = getRigth(i);

        // 最小节点的数组索引
        int smallest;

        // 比较得到父节点和子节点中的最小节点的数组索引
        if (left < size) {
            smallest = (array[i] <= array[left]) ? i : left;
        } else {
            smallest = i;
        }
        if (rigth < size) {
            smallest = (array[smallest] <= array[rigth]) ? smallest : rigth;
        }

        // 如果i与smallest不相等，表明父节点不满足小于等于子节点
        // 将父节点元素与最小子节点元素进行交换
        if (i != smallest) {
            int temp = array[i];
            array[i] = array[smallest];
            array[smallest] = temp;
            // 递归以防止交换后子节点不满足小根堆性质
            keepHeap(array, smallest, size);
        }
    }

}

大根堆与小根堆类似，这里就不再举大根堆的例子了。

三. 堆的构建

如果给定一个数组，里面元素并没有按照堆的性质进行放置，此时若需要根据数组中的元素构建一个堆，则应该让所有非叶子节点保持堆的性质（叶子节点总是满足堆的性质）。已经知道，一个堆的叶子节点总是在数组的最后部分，因此通过计算最后一个叶子节点（即数组中的最后一个位置）的父节点，就可以得到最后一个非叶子节点，该非叶子节点往前的所有节点都是非叶子节点，所以从最后一个非叶子节点开始往前，让每一个非叶子节点保持堆的性质，就可以构建一个堆。图示如下。

可以看到，叶子节点总是在数组的最后部分，并且通过计算最后一个叶子节点（数组索引为5）的父节点可以得到最后一个非叶子节点（数组索引为2）。

构建堆的JAVA代码如下所示。

public class HeapTry {

	// 省略其它

    public void createHeap(int[] array, int size) {
        // 计算最后一个非叶子节点的数组索引
        int lastParent = getParent(size - 1);

        // 从最后一个非叶子节点往前遍历每一个非叶子节点，并让每个非叶子节点保持堆性质
        for (int i = lastParent; i >= 0; i--) {
            keepHeap(array, i, size);
        }
    }
	
}

四. 堆的排序

因为小根堆的根节点的元素总是整个堆中的最小元素，大根堆的根节点的元素总是整个堆中的最大元素，因此可以基于堆的性质实现排序算法（通常小根堆实现降序排序，大根堆实现升序排序）。给定一个小根堆，且一共有n个元素，堆排序算法步骤如下。

1. 将小根堆的根节点与堆最后一个节点互换元素，然后执行步骤2；
2. 将前n-1个节点重新构建一个小根堆，并且令n=n-1，如果n为0则排序结束，否则执行步骤1。

图解流程如下所示。

堆排序（小根堆）的JAVA代码如下所示。

public class HeapTry {

	// 省略其它

    public void heapSort(int[] array, int size) {
        if (size == 0) {
            return;
        }
        while (size > 0) {
            int temp = array[0];
            array[0] = array[size - 1];
            array[size - 1] = temp;
            createHeap(array, size - 1);
            size = size - 1;
        }
    }
	
}

五. 插入和删除

如果要基于堆实现一个优先级队列，那么堆需要能够支持插入和删除元素。首先是元素的插入，以小根堆为例，将需要插入的元素添加到堆的最后以生成最后一个叶子节点，然后将其和父节点进行比较，如果其小于父节点则和父节点互换元素，然后继续和新位置对应的父节点进行比较，直到插入的元素大于等于父节点或者插入的元素成为根节点的元素为止。图示如下。

然后是元素的删除，同样以小根堆为例，每次删除都应该删除堆中的最小元素，即根节点的元素，然后将堆的最后一个叶子节点的元素填充到根节点，此时填充到根节点的元素可能破坏堆的性质，因此需要使用第二小节中的算法来恢复根节点的堆的性质。图示如下。

堆的插入和删除的JAVA代码如下所示。

public class HeapTry {

	// 省略其它

    public void insert(int[] array, int size, int i) {
        array[size] = i;
        int index = size;
        while (array[index] < array[getParent(index)]) {
            int temp = array[index];
            array[index] = array[getParent(index)];
            array[getParent(index)] = temp;
            index = (index - 1) / 2;
        }
    }

    public int pop(int[] array, int size) {
        if (size <= 0) {
            throw new IndexOutOfBoundsException();
        }
        int result = array[0];
        array[0] = array[size - 1];
        array[size - 1] = 0;
		keepHeap(array, 0, size - 1);
        return result;
    }

}

六. 实现优先级队列

本小节实现的优先级队列基于大根堆进行实现。要基于堆实现优先级队列，首先需要设计队列中的元素。先给出队列元素的JAVA代码，然后再进行说明。

public class PriObject<V> {

    private final V value;
    private final int priority;
    private final LocalDateTime dateTime;

    public PriObject(V value, int priority) {
        this.value = value;
        this.priority = priority;
        dateTime = LocalDateTime.now();
    }

    public V getValue() {
        return value;
    }

    public int getPriority() {
        return priority;
    }

    public LocalDateTime getDateTime() {
        return dateTime;
    }

    public int compareTo(PriObject<?> priObject) {
        if (priority > priObject.getPriority()) {
            return 1;
        } else if (priority < priObject.getPriority()) {
            return -1;
        } else {
            return -dateTime.compareTo(priObject.getDateTime());
        }
    }

}

队列元素类PriObject有三个属性，含义如下。

• value表示元素的值；
• priority表示元素的优先级，这里设定值越大优先级越高；
• dateTime表示元素的生成时间，当向优先级队列添加value时，优先级队列会将value封装成一个PriObject对象，dateTime则表示这个元素的生成时间，用于两个元素priority相等时判断优先级，设定时间越早越优先。

队列元素类PriObject还提供了compareTo()方法用于两个元素的比较，比较规则为：priority越大越优先，如果priority相等则dateTime越小越优先。

下面看一下优先级队列PriorityQueue的字段。

public class PriorityQueue<T> {

	// 堆数组初始容量为16
    private final int INITAIL_SIZE = 16;

	// 堆数组
    private PriObject<?>[] array = new PriObject<?>[INITAIL_SIZE];

	// 堆元素个数
    private final AtomicInteger size = new AtomicInteger(0);

	// 全局锁
    private final Lock lock = new ReentrantLock();
	
	// 省略其它
	
}

PriorityQueue的全局锁字段用于优先级队列在offer()元素和poll()元素时线程安全。

PriorityQueue内部也使用了方法来获取左右子节点和父节点的数组索引以及保持堆性质，如下所示。

public class PriorityQueue<T> {

	// 省略其它

    private int getLeft(int i) {
        return (i + 1) * 2 - 1;
    }

    private int getRight(int i) {
        return (i + 1) * 2;
    }

    private int getParent(int i) {
        if (i == 0) {
            return 0;
        }
        return (i - 1) / 2;
    }

    private void keepHeap(int i, int size) {
        int left = getLeft(i);
        int rigth = getRight(i);

        int max;

        if (left < size) {
            max = (array[i].compareTo(array[left]) > 0) ? i : left;
        } else {
            max = i;
        }
        if (rigth < size) {
            max = (array[max].compareTo(array[rigth]) > 0) ? max : rigth;
        }

        if (i != max) {
            PriObject<?> temp = array[i];
            array[i] = array[max];
            array[max] = temp;
            keepHeap(max, size);
        }
    }
	
	// 省略其它
	
}

下面看一下PriorityQueue的offer()方法。

public class PriorityQueue<T> {

	// 省略其它

    /**
     * 添加元素到队列，每一个元素会被封装为一个{@link PriObject}对象
     * 然后被添加到队列中。
     * 
     * @param value 队列存储的元素
     * @param priority 优先级
     * @return 添加成功返回true，失败返回false
     */
    public boolean offer(T value, int priority) {
        lock.lock();
        try {
            PriObject<T> pObj = new PriObject<>(value, priority);
            if (size.get() == array.length) {
                resize();
            }
            int index = size.get();
            array[index] = pObj;
            while (array[index].compareTo(array[getParent(index)]) > 0) {
                PriObject<?> temp = array[index];
                array[index] = array[getParent(index)];
                array[getParent(index)] = temp;
                index = getParent(index);
            }
            size.incrementAndGet();
        } catch (Exception e) {
            return false;
        } finally {
            lock.unlock();
        }
        return true;
    }
	
	private void resize() {
		int oldSize = array.length;
		int newSize = oldSize << 1;
		array = Arrays.copyOf(array, newSize);
    }
	
	// 省略其它
	
}

在向优先级队列添加一个元素时，如果元素个数在添加前已经等于堆数组长度，此时会触发扩容机制，并且扩容后容量为扩容前的一倍。

下面再看一下PriorityQueue的poll()方法。

public class PriorityQueue<T> {

	// 省略其它

    public T poll() {
        T value;
        lock.lock();
        try {
            int currentSize = size.get();
            if (currentSize == 0) {
                return null;
            }
            value = (T) array[0].getValue();
            array[0] = array[currentSize - 1];
            array[currentSize - 1] = null;
            keepHeap(0, currentSize - 1);
            size.decrementAndGet();
        } catch (Exception e) {
            value = null;
        } finally {
            lock.unlock();
        }
        return value;
    }

}

最后编写测试程序来验证实现的优先级队列的功能。测试代码如下所示。

public class PriorityQueueTest {

    private static final String EVENT_1 = "事件1";
    private static final String EVENT_2 = "事件2";
    private static final String EVENT_3 = "事件3";

    @Test
    public void 给定不同优先级事件_当调用offer方法_可以根据优先级正确获取元素() {
        Event event1 = new Event(EVENT_1);
        Event event2 = new Event(EVENT_2);
        Event event3 = new Event(EVENT_3);
        PriorityQueue<Event> priorityQueue = new PriorityQueue<>();

        priorityQueue.offer(event1, 10);
        priorityQueue.offer(event2, 20);
        priorityQueue.offer(event3, 5);

        assertThat(priorityQueue.poll().getEvent(), is(EVENT_2));
        assertThat(priorityQueue.poll().getEvent(), is(EVENT_1));
        assertThat(priorityQueue.poll().getEvent(), is(EVENT_3));
    }

    @Test
    public void 给定相同优先级事件_当调用offer方法_可以根据添加时间正确获取元素() {
        Event event1 = new Event(EVENT_1);
        Event event2 = new Event(EVENT_2);
        Event event3 = new Event(EVENT_3);
        PriorityQueue<Event> priorityQueue = new PriorityQueue<>();

        // 添加事件1，并睡眠10毫秒
        priorityQueue.offer(event1, 10);
        sleep10Millisecond();

        // 添加事件2，并睡眠10毫秒
        priorityQueue.offer(event2, 10);
        sleep10Millisecond();

        // 添加事件3
        priorityQueue.offer(event3, 10);

        assertThat(priorityQueue.poll().getEvent(), is(EVENT_1));
        assertThat(priorityQueue.poll().getEvent(), is(EVENT_2));
        assertThat(priorityQueue.poll().getEvent(), is(EVENT_3));
    }

    @Test
    public void 给定17个不同优先级的事件_当调用offer方法_触发扩容() {
        String eventString = "事件";
        PriorityQueue<Event> priorityQueue = new PriorityQueue<>();

        for (int i = 0; i < 17; i++) {
            Event event = new Event(eventString + i);
            priorityQueue.offer(event, i);
        }

        for (int i = 16; i >= 0; i--) {
            assertThat(priorityQueue.poll().getEvent(), is(eventString + i));
        }
    }

    @Test
    public void 给定17个相同优先级的事件_当调用offer方法_触发扩容() {
        String eventString = "事件";
        PriorityQueue<Event> priorityQueue = new PriorityQueue<>();

        for (int i = 0; i < 17; i++) {
            Event event = new Event(eventString + i);
            priorityQueue.offer(event, 10);
            sleep10Millisecond();
        }

        for (int i = 0; i < 17; i++) {
            assertThat(priorityQueue.poll().getEvent(), is(eventString + i));
        }
    }

    private void sleep10Millisecond() {
        try {
            Thread.sleep(10);
        } catch (InterruptedException e) {
            System.out.println(e.getMessage());
        }
    }

    private static class Event {
        private final String event;

        public Event(String event) {
            this.event = event;
        }

        public String getEvent() {
            return event;
        }
    }

}

测试了四个场景，分别是：priority不同，priority相同，触发扩容时priority不同和触发扩容时priority相同这四种场景下的优先级队列的功能实现。执行结果如下。

总结

堆是一颗完全二叉树，并且根据父节点总是大于等于子节点或者父节点总是小于等于子节点可将堆划分为大根堆和小根堆。堆中元素可以由一个数组来表示，并且可以根据某个节点在数组中的索引计算得到该节点的左右子节点和父节点在数组中的索引。基于堆可以实现排序算法，通常，大根堆实现升序排序，小根堆实现降序排序。基于堆还可以实现优先级队列，并且优先级队列的元素存储在一个堆数组中，堆数组在容量满时会进行扩容，因此可以将优先级队列看作是一个无界队列。