算法leetcode

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

字数统计: 2k字   |   阅读时长: 9分
, , , ,

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

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的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