跳至主要內容

1337. 矩阵中战斗力最弱的 K 行


1337. 矩阵中战斗力最弱的 K 行

🟢   🔖  数组 二分查找 矩阵 排序 堆(优先队列)  🔗 力扣open in new window LeetCodeopen in new window

题目

You are given an m x n binary matrix mat of 1's (representing soldiers) and 0's (representing civilians). The soldiers are positioned in front of the civilians. That is, all the 1's will appear to the left of all the 0's in each row.

A row i is weaker than a row j if one of the following is true:

  • The number of soldiers in row i is less than the number of soldiers in row j.
  • Both rows have the same number of soldiers and i < j.

Return the indices of thek weakest rows in the matrix ordered from weakest to strongest.

Example 1:

Input:

mat =
[[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]],

k = 3

Output: [2,0,3]

Explanation:

The number of soldiers in each row is:

  • Row 0: 2
  • Row 1: 4
  • Row 2: 1
  • Row 3: 2
  • Row 4: 5

The rows ordered from weakest to strongest are [2,0,3,1,4].

Example 2:

Input:

mat =
[[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]],

k = 2

Output: [0,2]

Explanation:

The number of soldiers in each row is:

  • Row 0: 1
  • Row 1: 4
  • Row 2: 1
  • Row 3: 1

The rows ordered from weakest to strongest are [0,2,3,1].

Constraints:

  • m == mat.length
  • n == mat[i].length
  • 2 <= n, m <= 100
  • 1 <= k <= m
  • matrix[i][j] is either 0 or 1.

题目大意

给你一个大小为 m * n 的矩阵 mat,矩阵由若干军人和平民组成,分别用 1 和 0 表示。

请你返回矩阵中战斗力最弱的 k 行的索引,按从最弱到最强排序。

如果第 i 行的军人数量少于第 j 行,或者两行军人数量相同但 i 小于 j ,那么我们认为第 i 行的战斗力比第 j 行弱。

军人 总是 排在一行中的靠前位置,也就是说 1 总是出现在 0 之前。

示例 1:

输入:

mat =
[[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]],

k = 3

输出:[2,0,3]

解释:

每行中的军人数目:

  • 行 0 -> 2
  • 行 1 -> 4
  • 行 2 -> 1
  • 行 3 -> 2
  • 行 4 -> 5

从最弱到最强对这些行排序后得到 [2,0,3,1,4]

示例 2:

输入:

mat =
[[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]],

k = 2

输出:[0,2]

解释:

每行中的军人数目:

  • 行 0 -> 1
  • 行 1 -> 4
  • 行 2 -> 1
  • 行 3 -> 1

从最弱到最强对这些行排序后得到 [0,2,3,1]

提示:

  • m == mat.length
  • n == mat[i].length
  • 2 <= n, m <= 100
  • 1 <= k <= m
  • matrix[i][j] 不是 0 就是 1

解题思路

  1. 矩阵行的强度计算:

    • 使用 map 将每一行的士兵数量及其索引存储为 [strength, index]
    • 每一行的强度为该行中值为 1 的数量,可以通过 reduce 或二分查找(针对排序行)快速计算。
  2. 构建最小堆:
    定义一个 MinHeap 类实现最小堆操作,提供 insertpop 方法。

  3. 堆优化:
    使用最小堆(Min-Heap)存储计算得到的 [strength, index] 数据,其中堆的优先级为:

    • 优先按士兵数量升序排列;
    • 如果士兵数量相同,按行号升序排列。
  4. 提取最弱的行:

    • 使用堆连续 kpop 操作,得到最弱的 k 行,返回对应的 index 数组。

复杂度分析

  • 时间复杂度O(m * n + k log m)
    • 计算强度: O(m * n),其中 m 是矩阵的行数,n 是列数。
    • 堆构建: O(m)
    • 提取最弱行: O(k log m)
    • 整体复杂度为 O(m * n + k log m)
  • 空间复杂度O(m),用于存储堆的数据。

代码

/**
 * @param {number[][]} mat
 * @param {number} k
 * @return {number[]}
 */
var kWeakestRows = function (mat, k) {
	// 自定义排序规则
	const priority = (a, b) => {
		if (a[0] === b[0]) return a[1] < b[1];
		return a[0] < b[0];
	};

	// 计算每一行的士兵数量及索引
	const arr = mat.map((row, i) => [row.reduce((a, b) => a + b, 0), i]);

	// 构建最小堆
	let minHeap = new MinHeap(arr, priority);

	// 提取最弱的 k 行
	let res = [];
	while (k-- > 0) {
		res.push(minHeap.pop()[1]);
	}

	return res;
};

// 最小堆实现
class MinHeap {
	constructor(arr = [], priority = (a, b) => a < b) {
		this.heap = arr;
		this.priority = priority;
		for (let i = ((this.heap.length / 2) | 0) - 1; i >= 0; i--) {
			this.heapifyDown(i);
		}
	}

	// 插入元素
	insert(item) {
		this.heap.push(item);
		this.heapifyUp(this.heap.length - 1);
	}

	// 移除并返回堆顶
	pop() {
		if (this.heap.length === 0) return null;
		const top = this.heap[0];
		const last = this.heap.pop();
		if (this.heap.length) {
			this.heap[0] = last;
			this.heapifyDown(0);
		}
		return top;
	}

	heapifyUp(i) {
		while (i > 0) {
			const parent = ((i - 1) / 2) | 0;
			if (this.priority(this.heap[i], this.heap[parent])) {
				[this.heap[i], this.heap[parent]] = [this.heap[parent], this.heap[i]];
				i = parent;
			} else {
				break;
			}
		}
	}

	heapifyDown(i) {
		const left = i * 2 + 1,
			right = i * 2 + 2;
		let min = i;
		if (
			left < this.heap.length &&
			this.priority(this.heap[left], this.heap[min])
		) {
			min = left;
		}
		if (
			right < this.heap.length &&
			this.priority(this.heap[right], this.heap[min])
		) {
			min = right;
		}
		if (min !== i) {
			[this.heap[i], this.heap[min]] = [this.heap[min], this.heap[i]];
			this.heapifyDown(min);
		}
	}
}