11 | 面试题：返回数据流中的第K大元素-算法面试通关 50 讲

当前位置:　首页>> 技术小册>> 算法面试通关 50 讲

### 第十一章 面试题：返回数据流中的第K大元素

在算法面试中，处理数据流问题是一项常见的挑战，它要求算法能够高效地处理连续到达的数据项，同时保持较低的内存占用。其中，“返回数据流中的第K大元素”是一个极具代表性的问题，它不仅考察了候选人的数据结构选择能力，还检验了其对排序、堆等高级算法概念的理解和应用。本章将深入剖析这一问题，从问题分析、解决方案设计到代码实现，全方位解析如何高效地解决这一面试难题。

#### 11.1 问题描述与分析

**问题描述**：给定一个数据流（可能包含重复元素），以及一个整数K，设计一个算法来找到数据流中第K大的元素。数据流以连续的方式到达，且元素的数量可能非常大，因此不能将所有元素存储在内存中。算法需要支持两个操作：`add(val)`，向数据流中添加一个新元素`val`；`findKthLargest(k)`，返回数据流中当前第K大的元素。

**问题分析**：
- **实时性**：由于数据流是连续到达的，算法需要能够实时或近乎实时地处理新元素，并更新第K大的元素。
- **内存限制**：由于元素数量可能非常大，算法不能简单地将所有元素存储在数组中，而是需要采用一种空间复杂度较低的数据结构。
- **动态性**：随着新元素的加入，数据流中的第K大元素可能会发生变化，算法需要能够动态地适应这种变化。

#### 11.2 解决方案设计

为了解决这个问题，我们可以考虑使用以下几种数据结构：

1. **优先队列（最小堆）**：
   - 优先队列（或称为堆）是一种特殊的完全二叉树结构，它能够在O(log n)时间复杂度内完成插入和删除最小（或最大）元素的操作。
   - 对于这个问题，我们可以使用一个大小为K的最小堆来维护数据流中当前最大的K个元素。当新元素加入时，如果堆未满（元素数量小于K），则直接将其加入堆中；如果堆已满，则先与新元素比较，如果新元素大于堆顶元素（即当前第K大的元素），则替换堆顶元素并重新调整堆结构。
   - 查找第K大元素时，直接返回堆顶元素即可。

2. **有序集合（如平衡二叉搜索树）**：
   - 虽然堆是一种有效的解决方案，但在某些情况下，使用有序集合（如AVL树、红黑树等）可能提供更高的灵活性。
   - 有序集合能够保持元素的有序性，并支持快速的插入、删除和查找操作。但是，与堆相比，其维护成本可能更高。
   - 对于这个问题，我们可以使用有序集合来维护数据流中所有的元素，但在每次插入新元素后，都需要确保集合中的元素保持有序，并快速定位到第K大的元素。这通常涉及到更复杂的树结构调整。

3. **分治策略（基于快速选择算法）**：
   - 如果允许一定的随机性，并且数据流不是实时要求更新第K大元素，可以考虑使用分治策略，如快速选择算法。
   - 快速选择算法是快速排序的一个变种，它可以在平均O(n)时间复杂度内找到一个未排序数组中第K小的元素。然而，这种方法需要遍历整个数组，对于数据流问题来说，每次添加新元素时都需要重新执行快速选择，效率较低。
   - 但如果数据流不是实时性要求极高，且可以接受一定的误差范围，快速选择算法可以作为一种备选方案。

综合考虑，对于“返回数据流中的第K大元素”这一问题，使用大小为K的最小堆是最直接且高效的解决方案。

#### 11.3 代码实现

以下是使用Python实现的一个简单示例，使用最小堆（通过`heapq`模块）来解决问题：

```python
import heapq

class KthLargest:
    def __init__(self, k: int, nums: List[int] = []):
        self.k = k
        self.heap = []
        # 初始化时将前k个元素加入堆中
        for num in nums[:k]:
            heapq.heappush(self.heap, num)
        # 如果初始数据超过k个，则删除堆中最小（第k+1大）的元素
        while len(self.heap) > k:
            heapq.heappop(self.heap)

def add(self, val: int) -> None:
        # 如果新元素大于堆顶元素，则替换并重新调整堆
        if len(self.heap) < self.k or val > self.heap[0]:
            if len(self.heap) == self.k:
                heapq.heappop(self.heap)
            heapq.heappush(self.heap, val)

def findKthLargest(self, k: int) -> int:
        # 这里k的值应与构造函数中的k一致，但为了接口灵活性仍保留参数
        if k != self.k:
            raise ValueError("The specified k does not match the initialized k.")
        return self.heap[0]

# 示例用法
kthLargest = KthLargest(3, [4, 5, 8, 2])
print(kthLargest.findKthLargest(3))  # 输出 4
kthLargest.add(3)
print(kthLargest.findKthLargest(3))  # 输出 4
kthLargest.add(6)
print(kthLargest.findKthLargest(3))  # 输出 5
kthLargest.add(7)
print(kthLargest.findKthLargest(3))  # 输出 5
```

#### 11.4 复杂度分析

- **时间复杂度**：
  - `add`操作：O(log K)，因为每次插入新元素时，可能需要将新元素与堆中的元素进行比较，并重新调整堆结构。
  - `findKthLargest`操作：O(1)，因为第K大的元素始终位于堆顶。
- **空间复杂度**：O(K)，因为我们需要维护一个大小为K的堆来存储当前最大的K个元素。

#### 11.5 总结

本章通过“返回数据流中的第K大元素”这一面试题，深入探讨了数据流处理问题的解决思路，重点介绍了使用最小堆作为数据结构的解决方案，并提供了详细的代码实现和复杂度分析。希望读者能够从中掌握处理数据流问题的基本方法，并在未来的面试和工作中灵活运用。同时，也鼓励读者尝试使用其他数据结构（如有序集合）或算法（如快速选择）来解决该问题，以加深对算法和数据结构的理解。

该分类下的相关小册推荐：

编程之道-算法面试(上)

数据结构与算法(上)

业务开发实用算法精讲

数据结构与算法(中)

数据结构与算法之美

数据结构与算法(下)

编程之道-算法面试(下)