【搜索难题】力扣1439-有序矩阵中的第k个最小数组和

  |  

摘要: 一个比较难得搜索题

【对算法,数学,计算机感兴趣的同学,欢迎关注我哈,阅读更多原创文章】
我的网站:潮汐朝夕的生活实验室
我的公众号:算法题刷刷
我的知乎:潮汐朝夕
我的github:FennelDumplings
我的leetcode:FennelDumplings


各位好,本文我们来看一个搜索题,这是力扣第 187 场周赛 D 题。首先本题可以抽象成 TopK 问题,于是基于堆的算法可以解决,但是队中维护的元素比较复杂。此外还可以对结果值域二分,判定过程用 DFS 算法,可以对搜索过程进行一些优化,也是比较常规的做法。


$1 题目

给你一个 m * n 的矩阵 mat,以及一个整数 k ,矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。

提示:

1
2
3
4
5
6
m == mat.length
n == mat.length[i]
1 <= m, n <= 40
1 <= k <= min(200, n ^ m)
1 <= mat[i][j] <= 5000
mat[i] 是一个非递减数组

示例 1:
输入:mat = [[1,3,11],[2,4,6]], k = 5
输出:7
解释:从每一行中选出一个元素,前 k 个和最小的数组分别是:
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。

示例 2:
输入:mat = [[1,3,11],[2,4,6]], k = 9
输出:17

示例 3:
输入:mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出:9
解释:从每一行中选出一个元素,前 k 个和最小的数组分别是:
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。

示例 4:
输入:mat = [[1,1,10],[2,2,9]], k = 7
输出:12

$2 题解

算法1:堆维护 topK

定义状态,其中有各行选中的下标的信息,以及这些元素的和的信息。

1
2
3
4
5
6
struct Item
{
vector<int> idxs;
int sum;
Item(const vector<int>& idxs, int sum):idxs(idxs),sum(sum){}
};

在堆中维护这些状态,并且按照和从小到大排序。每从堆中弹出状态都是从全局来看和最小的。该状态的下标列表中,每个下标自增1,都会产生下一个状态。将这些状态压堆。

压堆之前需要用哈希表维护一下去重,c++ 实现需要写 MyHash

我们需要对 vector<int> 定义哈希函数,在 参考字符串哈希定义数组哈希(数组的同构) 中我们处理过这个问题,代码模板如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
const int MOD = 1610612741;
const int p = 201326611;
using ll = long long;

struct MyHash
{
int operator()(const vector<int>& idxs) const
{
int ans = 0;
for(int i: idxs)
{
ans = (ans * (ll)p) % MOD;
ans = (ans + (ll)i) % MOD;
}
return ans;
}
};

unordered_set<vector<int>, MyHash> setting;

代码 (C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
struct Item
{
vector<int> idxs;
int sum;
Item(const vector<int>& idxs, int sum):idxs(idxs),sum(sum){}
};

struct HeapCmp
{
bool operator()(const Item& i1, const Item& i2) const
{
return i1.sum > i2.sum;
}
};

const int MOD = 1610612741;
const int p = 201326611;
using ll = long long;

struct MyHash
{
int operator()(const vector<int>& idxs) const
{
int ans = 0;
for(int i: idxs)
{
ans = (ans * (ll)p) % MOD;
ans = (ans + (ll)i) % MOD;
}
return ans;
}
};

class Solution {
public:
int kthSmallest(vector<vector<int>>& mat, int k) {
int n = mat.size(), m = mat[0].size();
int sum = 0;
for(int i = 0; i < n; ++i)
sum += mat[i][0];
vector<int> idxs(n, 0);
priority_queue<Item, vector<Item>, HeapCmp> pq;
pq.push(Item(idxs, sum));
unordered_set<vector<int>, MyHash> setting;
setting.insert(idxs);
while(!pq.empty() && k > 0)
{
vector<int> nxt_idxs = pq.top().idxs;
int cur_sum = pq.top().sum;
pq.pop();
--k;
if(k == 0)
return cur_sum;
for(int i = 0; i < n; ++i)
{
int j = nxt_idxs[i];
if(j + 1 == m) continue;
int nxt_sum = cur_sum - mat[i][j] + mat[i][j + 1];
nxt_idxs[i] = j + 1;
if(setting.count(nxt_idxs) == 0)
{
setting.insert(nxt_idxs);
pq.push(Item(nxt_idxs, nxt_sum));
}
nxt_idxs[i] = j;
}
}
return -1;
}
};

算法2: 值域二分 + DFS

1
2
3
4
5
6
7
8
9
10
11
最大和 left = sum(mat[i][0]), 
最大和 right = sum(mat[i][m-1])

二分可能的和 mid = (left + right + 1) / 2
int cnt = check(mid) 和小于 mid 的数组的个数
if(cnt >= K)
mid 取大了
right = mid - 1
else
mid 可能最终会是答案
left = mid

check 的实现:和小于 mid 的数组个数

1
2
3
4
5
6
7
bool dfs(vector<int>& idxs, int& cnt, const int& mid);

对于当前数组 idxs:
如果其和 >= mid,直接返回
如果其和 < mid,++cnt,然后继续搜索 idxs 各个位置自增 1 的新状态

同样需要实现哈希表维护去重

代码 (C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class Solution {
public:
int kthSmallest(vector<vector<int>>& mat, int k) {
start_sum = 0;
int n = mat.size(), m = mat[0].size();
int right = 0;
for(int i = 0; i < n; ++i)
{
right += mat[i][m - 1];
start_sum += mat[i][0];
}
int left = start_sum;
while(left < right)
{
int mid = (left + right + 1) / 2;
int cnt = check(mat, mid, k);
if(cnt >= k)
right = mid - 1;
else
left = mid;
}
return left;
}

private:
int start_sum;
vector<int> start;
unordered_set<vector<int>, MyHash> setting;

int check(const vector<vector<int>>& mat, int mid, int k)
{
int ans = 0;
int n = mat.size();
start.assign(n, 0);
setting.clear();
setting.insert(start);
dfs(mat, start, start_sum, ans, mid, k);
return ans;
}

void dfs(const vector<vector<int>>& mat, vector<int>& idxs, int sum, int& cnt, const int mid, const int k)
{
if(sum >= mid)
return;
++cnt;
if(cnt > k)
return;
int n = mat.size(), m = mat[0].size();
for(int i = 0; i < n; ++i)
{
int j = idxs[i];
if(j + 1 < m)
{
idxs[i] = j + 1;
if(setting.count(idxs) == 0)
{
setting.insert(idxs);
int nxt_sum = sum + mat[i][j + 1] - mat[i][j];
dfs(mat, idxs, nxt_sum, cnt, mid, k);
}
idxs[i] = j;
}
}
}
};

优化:搜索过程优化

dfs 参数带着 i, sum, cnt, mid, k,当前 dfs 的这一枝表示的是前 i 行 ([0..i-1]) 做出特定选择的情况下,尝试推进 i 这一行的 j,1 向 m - 1 推,直至推不动。

前面 [0..i-1] 行的具体选择不用关心,只需要知道该选择下对应的当前和,此数据由 dfs 参数 sum 持有。

在该行每合法推进一次 j,都将此选择结合前 i 行的选择形成第 [0..i] 的一次合法选择然后搜在此基础上推进下一行的结果。

在每一枝的搜索之前,可以做写剪枝操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class Solution {
public:
int kthSmallest(vector<vector<int>>& mat, int k) {
start_sum = 0;
int n = mat.size(), m = mat[0].size();
int right = 0;
for(int i = 0; i < n; ++i)
{
right += mat[i][m - 1];
start_sum += mat[i][0];
}
int left = start_sum;
while(left < right)
{
int mid = (left + right + 1) / 2;
int cnt = check(mat, mid, k);
if(cnt >= k)
right = mid - 1;
else
left = mid;
}
return left;
}

private:
int start_sum;

int check(const vector<vector<int>>& mat, int mid, int k)
{
// ans: 于 mid 的个数
// ans >= k,可以早停
int ans = 1;
dfs(mat, 0, start_sum, ans, mid, k);
return ans;
}

// dfs 参数迭代 i,内部迭代 j (i, j) 对应
void dfs(const vector<vector<int>>& mat, int i, int sum, int& cnt, const int mid, const int k)
{
// 推进 mat 第 i 行的 j
// sum 是在 [0..i-1] 行已经选了某些点的情况下的和
// 此时 [i..n-1] 行仍然是用的 mat[i..n-1][0],即没有在前 i 行做 sum 代表的的选择下选
int n = mat.size(), m = mat[0].size();
// 剪枝
if(sum >= mid || cnt >= k || i == n)
return;
// sum < mid, cnt < k, i < n
// sum 是个合法选择
// sum 代表的选择在调用方已经计了 cnt
// 在 sum 选择下先搜索下一行
dfs(mat, i + 1, sum, cnt, mid, k);
// 在回溯阶段推进本行的 j
int j = 1;
sum -= mat[i][0];
while(j < m && sum + mat[i][j] < mid)
{
++cnt;
dfs(mat, i + 1, sum + mat[i][j], cnt, mid, k);
++j;
}
}
};

Share