堆与优先队列实现 | 程序员小乐

算法堆优先队列

算法

发布日期: 2024-09-09

作者: 小乐

文章字数: 909

阅读时长: 4 分

阅读次数:

一、堆的定义

1.1 堆的性质

堆是完全二叉树，满足堆序性质：

大顶堆：父节点 ≥ 子节点
小顶堆：父节点 ≤ 子节点

1.2 数组存储

利用完全二叉树特性，用数组存储：

索引i的节点：
- 父节点：(i - 1) / 2
- 左子节点：2i + 1
- 右子节点：2i + 2

        10
       /  \
      8    7
     / \   / \
    5   6 4   3

数组: [10, 8, 7, 5, 6, 4, 3]

二、堆的核心操作

2.1 上浮(Sift Up)

新元素插入末尾，与父节点比较，不满足堆序则交换。

void siftUp(int[] heap, int i) {
    while (i > 0) {
        int parent = (i - 1) / 2;
        if (heap[parent] >= heap[i]) break;  // 大顶堆
        swap(heap, i, parent);
        i = parent;
    }
}

时间复杂度：( O(\log n) )

2.2 下沉(Sift Down)

堆顶元素移除后，末尾元素放堆顶，与子节点比较交换。

void siftDown(int[] heap, int n, int i) {
    while (true) {
        int largest = i;
        int left = 2 * i + 1;
        int right = 2 * i + 2;
        
        if (left < n && heap[left] > heap[largest]) largest = left;
        if (right < n && heap[right] > heap[largest]) largest = right;
        
        if (largest == i) break;
        swap(heap, i, largest);
        i = largest;
    }
}

时间复杂度：( O(\log n) )

2.3 建堆(Heapify)

从最后一个非叶子节点开始下沉：

void buildHeap(int[] arr) {
    int n = arr.length;
    for (int i = n / 2 - 1; i >= 0; i--) {
        siftDown(arr, n, i);
    }
}

时间复杂度：( O(n) )（不是 ( O(n \log n) )）

证明：高度为 ( h ) 的节点最多有 ( \lceil n / 2^{h+1} \rceil ) 个，每个下沉 ( O(h) )。
[ \sum_{h=0}^{\log n} \frac{n}{2^{h+1}} \cdot h = O(n) ]

三、堆排序

void heapSort(int[] arr) {
    buildHeap(arr);  // 建大顶堆
    
    for (int i = arr.length - 1; i > 0; i--) {
        swap(arr, 0, i);      // 堆顶放末尾
        siftDown(arr, i, 0);  // 剩余元素重新堆化
    }
}

时间复杂度：( O(n \log n) )
空间复杂度：( O(1) )

四、优先队列

4.1 接口定义

interface PriorityQueue<T> {
    void offer(T element);   // 入队 O(log n)
    T poll();                // 出队（优先级最高）O(log n)
    T peek();                // 查看队首 O(1)
}

4.2 Java PriorityQueue

// 小顶堆（默认）
PriorityQueue<Integer> pq = new PriorityQueue<>();

// 大顶堆
PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());

// 自定义比较器
PriorityQueue<Task> pq = new PriorityQueue<>(
    Comparator.comparingInt(t -> t.priority)
);

4.3 底层实现

Java PriorityQueue 使用数组实现的二叉堆：

transient Object[] queue;  // 堆数组
int size;                   // 元素数量
Comparator<? super E> comparator;

五、经典应用场景

5.1 Top K问题

求数组中第K大元素：

public int findKthLargest(int[] nums, int k) {
    // 小顶堆，维护K个最大元素
    PriorityQueue<Integer> pq = new PriorityQueue<>(k);
    for (int num : nums) {
        pq.offer(num);
        if (pq.size() > k) pq.poll();
    }
    return pq.peek();
}

时间复杂度：( O(n \log k) )

5.2 合并K个有序数组

public int[] mergeKLists(int[][] lists) {
    PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
    
    // 每个数组的当前元素入堆
    for (int i = 0; i < lists.length; i++) {
        if (lists[i].length > 0) {
            pq.offer(new int[]{lists[i][0], i, 0});
        }
    }
    
    List<Integer> res = new ArrayList<>();
    while (!pq.isEmpty()) {
        int[] curr = pq.poll();
        res.add(curr[0]);
        int arrIdx = curr[1], elemIdx = curr[2];
        if (elemIdx + 1 < lists[arrIdx].length) {
            pq.offer(new int[]{lists[arrIdx][elemIdx + 1], arrIdx, elemIdx + 1});
        }
    }
    return res.stream().mapToInt(i -> i).toArray();
}