二分查找与堆排序解决矩阵TopK问题

1. 题目解析与核心思路

1337题"矩阵中战斗力最弱的K行"是一道典型的矩阵处理与排序结合的算法题。题目给出一个由0和1组成的m×n矩阵，其中1总是出现在0的左侧（即每行都是非递增序列）。要求根据每行中1的数量（即"战斗力"）进行排序，返回战斗力最弱的K行的行号。

这个问题的核心在于两个关键操作：

1. 计算每行中1的数量（即战斗力）
2. 根据战斗力对行进行排序，并处理战斗力相同的情况

在实际解题中，我采用了二分查找统计每行1的数量，然后使用优先队列（堆）来处理排序和取前K个元素的操作。这种方法的时间复杂度为O(m log n + m log k)，相比暴力解法有显著优化。

2. 二分查找统计战斗力

2.1 为什么选择二分查找

由于题目保证每行都是非递增序列（即1后面不会出现0），这为二分查找提供了理想条件。对于每行，我们可以：

1. 找到最后一个1的位置
2. 该位置+1即为该行1的总数

相比线性扫描O(n)的时间复杂度，二分查找可以将单行的统计时间降至O(log n)。这在处理大规模矩阵时优势明显。

2.2 二分查找实现细节

python复制def count_ones(row):
    left, right = 0, len(row)
    while left < right:
        mid = (left + right) // 2
        if row[mid] == 1:
            left = mid + 1
        else:
            right = mid
    return left

这个实现有几个关键点：

1. 使用左闭右开区间[left, right)
2. 当row[mid]为1时，说明最后一个1可能在右侧，因此left=mid+1
3. 当row[mid]为0时，说明最后一个1可能在左侧或mid就是第一个0

注意：二分查找的边界条件容易出错，建议在纸上模拟几个测试用例验证逻辑正确性。

3. 优先队列处理TopK问题

3.1 堆的选择与使用

统计完所有行的战斗力后，我们需要找出战斗力最弱的K行。这是一个典型的TopK问题，优先队列（堆）是最佳选择：

1. 使用最大堆（大顶堆）维护当前找到的K个最弱行
2. 当新行的战斗力小于堆顶时，替换堆顶元素
3. 最终堆中元素即为所需的K个最弱行

Python中的heapq模块默认实现最小堆，因此我们需要存储负值来模拟最大堆：

python复制import heapq

def kWeakestRows(mat, k):
    counts = [(count_ones(row), i) for i, row in enumerate(mat)]
    heap = []
    for cnt, idx in counts:
        if len(heap) < k:
            heapq.heappush(heap, (-cnt, -idx))
        else:
            if -cnt > heap[0][0] or (-cnt == heap[0][0] and -idx > heap[0][1]):
                heapq.heappushpop(heap, (-cnt, -idx))
    return [-idx for cnt, idx in sorted(heap, reverse=True)]

3.2 处理相同战斗力的特殊情况

当两行战斗力相同时，题目要求按行号从小到大排序。这在堆中需要特殊处理：

1. 比较战斗力
2. 战斗力相同时比较行号
3. 由于使用最大堆，行号也需要取负

4. 完整代码实现与优化

4.1 最终优化版本

结合上述思路，完整的优化实现如下：

python复制import heapq

def kWeakestRows(mat, k):
    def count_ones(row):
        left, right = 0, len(row)
        while left < right:
            mid = (left + right) // 2
            if row[mid] == 1:
                left = mid + 1
            else:
                right = mid
        return left
    
    counts = [(count_ones(row), i) for i, row in enumerate(mat)]
    heap = []
    
    for cnt, idx in counts:
        if len(heap) < k:
            heapq.heappush(heap, (-cnt, -idx))
        else:
            if -cnt > heap[0][0] or (-cnt == heap[0][0] and -idx > heap[0][1]):
                heapq.heappushpop(heap, (-cnt, -idx))
    
    result = []
    while heap:
        cnt, idx = heapq.heappop(heap)
        result.append(-idx)
    
    return result[::-1]

4.2 复杂度分析

1. 统计战斗力：O(m log n) - m行每行二分查找
2. 构建堆：O(m log k) - 最多m次堆操作，每次O(log k)
3. 总复杂度：O(m(log n + log k)) = O(m log(nk))

空间复杂度：O(k) - 堆的大小最多为k

5. 边界条件与测试用例

5.1 常见边界情况

1. K等于矩阵行数：返回所有行
2. K=1：只返回最弱的一行
3. 所有行战斗力相同：按行号排序
4. 矩阵为空：返回空列表
5. 矩阵只有一列：所有行战斗力为0或1

5.2 测试用例示例

python复制test_cases = [
    # 常规情况
    ([[1,1,0,0,0],
      [1,1,1,1,0],
      [1,0,0,0,0],
      [1,1,0,0,0],
      [1,1,1,1,1]], 3, [2,0,3]),
    
    # 所有行战斗力相同
    ([[1,0],
      [1,0]], 1, [0]),
    
    # K等于行数
    ([[1,1,0],
      [1,0,0]], 2, [1,0]),
    
    # 边缘情况
    ([], 0, []),
    ([[0,0],[0,0]], 1, [0])
]

6. 性能优化与替代方案

6.1 垂直扫描法

另一种思路是垂直扫描矩阵列，当发现某列有0时，检查该行是否全为1。这种方法最优情况下可以达到O(m+n)时间复杂度：

python复制def kWeakestRows(mat, k):
    m, n = len(mat), len(mat[0])
    result = []
    visited = set()
    
    for j in range(n):
        for i in range(m):
            if i not in visited and mat[i][j] == 0:
                result.append(i)
                visited.add(i)
                if len(result) == k:
                    return result
    
    # 处理全1行
    for i in range(m):
        if i not in visited:
            result.append(i)
            if len(result) == k:
                break
    return result

6.2 性能对比

在实际LeetCode测试中，二分+堆的方法运行时间稳定在100ms以内，是最可靠的解决方案。

7. 实际应用与扩展

7.1 现实世界中的应用

这种矩阵处理技术可以应用于：

1. 用户行为分析（如点击流数据）
2. 图像处理中的特征提取
3. 推荐系统中的用户-物品矩阵
4. 社交网络中的连接矩阵分析

7.2 题目变种与扩展

1. 如果矩阵不是严格非递增怎么办？

   • 需要先对每行进行排序，增加O(mn log n)预处理时间

2. 如果要求最强的K行？

   • 只需修改堆的比较逻辑

3. 如果矩阵非常大无法放入内存？

   • 需要使用外部排序和流式处理技术

4. 如果要求动态更新和查询？

   • 需要设计更复杂的数据结构如二叉索引树

在实际编码中，我发现这类矩阵问题的关键在于充分利用题目给出的特殊条件（如本题的非递增性质）。这往往能带来算法上的优化空间。同时，处理TopK问题时，堆通常是首选数据结构，但要注意处理相同值的情况。