剑指 Offer 04. 二维数组中的查找题解

题目链接

题目大意：给出一个 $n\times m$ 的整数数组，数组中数从左向右、从上到下递增，判断该二维数组中是否含有某个整数。

最直接的思路是遍历整个二维数组，但因为是从上到下、从左到右递增，所以如果目标整数小于当前行最左边的值，则不用再往下遍历；如果目标整数大于最右边的值，则遍历下一行。

不过第一次代码被卡了，发现是因为当数组为空时会发生段错误，所以增加了判断数组是否为空。

代码：

#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    bool findNumberIn2DArray(vector<vector<int>>& matrix, int target) {
        for (auto i = matrix.begin(); i != matrix.end(); i++) {
            if (i->begin() == i->end()) {
                return false;
            }
            if (target < *(i->begin())) {
                return false;
            }
            if (target > *(i->end() - 1)) {
                continue;
            }
            for (auto j = i->begin(); j != i->end(); j++) {
                if (*j == target) {
                    return true;
                }
            }
        }
        return false;
    }
};

int main(void) {
    Solution s;

    vector<int> a{1, 4, 7, 11, 15};
    vector<int> b{2, 5, 8, 12, 19};
    vector<int> c{3, 6, 9, 16, 22};
    vector<int> d{10, 13, 14, 17, 24};
    vector<int> e{18, 21, 23, 26, 30};

    vector<vector<int>> v;
    v.push_back(a);
    v.push_back(b);
    v.push_back(c);
    v.push_back(d);
    v.push_back(e);

    cout << s.findNumberIn2DArray(v, 5) << endl;
    cout << s.findNumberIn2DArray(v, 20) << endl;

    return 0;
}

我接下来优化的思路是把每行的遍历转换为二分查找，时间复杂度从 $O(nm)$ 变为 $O(nlo{g}_m)$。

代码：

class Solution {
public:
    bool findNumberIn2DArray(vector<vector<int>>& matrix, int target) {
        for (auto i = matrix.begin(); i != matrix.end(); i++) {
            if (i->begin() == i->end()) {
                return false;
            }
            if (target < *(i->begin())) {
                return false;
            }
            if (target > *(i->end() - 1)) {
                continue;
            }
            
            int left = 0, right = i->size() - 1, middle;
            if (target == (*i)[left] || target == (*i)[right]) {
                return true;
            }
            while (right - left > 1) {
                middle = (left + right) / 2;
                if (target == (*i)[middle]) {
                    return true;
                } else if (target > (*i)[middle]) {
                    left = middle;
                } else {
                    right = middle;
                }
            }
        }
        return false;
    }
};

后来我看了题解，发现可以把时间复杂度优化到 $O(n+m)$。其思路是将二维数组旋转 $45^{\circ }$，转换为一个类似二叉搜索树的图。相当于从右上角开始搜索，如果目标整数大于当前数，则列的值加一；如果目标整数小于当前数，则行的值减一。

class Solution {
public:
    bool findNumberIn2DArray(vector<vector<int>>& matrix, int target) {
        if (matrix.size() == 0 || matrix[0].size() == 0) {
            return false;
        }

        int i = 0, j = matrix[0].size() - 1;
        while (i < matrix.size() && j >= 0) {
            if (target == matrix[i][j]) {
                return true;
            } else if (target < matrix[i][j]) {
                j--;
            } else {
                i++;
            }
        }
        return false;
    }
};