LeetCode #378 有序矩阵中第 K 小的元素

本贴最后更新于 1320 天前,其中的信息可能已经时移俗易

378.png

Problem Description

给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。

note

你可以假设 k 的值永远是有效的,1 ≤ k ≤ n ^ 2

e.g.

  • 示例:

    matrix = [
       [1,  5,  9],
       [10, 11, 13],
       [12, 13, 15]
    ],
    k = 8,
    
    返回 13。
    

Solution

有一说一,被二维数组中的查找这道题整魔怔了。该题目的一种解法是将该矩阵模拟成一个二叉搜索树,利用其性质来做遍历方向的决策从而快速找到 target。因为是模拟二叉搜索树,所以查找的时间复杂度是 O(log(n ^ 2) )害挺快的。而这题也是同样的矩阵,但求的是第 k 个值(排序好的第 K 位)。

刚好晚上看了麻省理工公开课的其中一节课(二叉搜索树和快排)

认识到了二者相似性和奇妙之处。二叉搜索树的构建过程和快排的排序过程,二者在最坏的情况下时间复杂度都是 O(n ^ 2) ,平均时间复杂度都是 O(nlog(n) ),原理十分类似。BST 在最差的情况是什么呢?就是像链表一样的情况,只有左节点或者右节点,这样它的高度就是节点数量(N),而在满二叉树下它的高度是(logN)。所以 N 个元素在构建 BST 的过程中,会先查找它应该在的位置(logN),N 个元素就是 O(nlog(n) )。但是在有序的数列(正、逆序)中会出现很糟糕的情况,这点和快排的 pivot 的选择很像。

像链表一样

所以这道题我就在想可不可以利用 BST 的中序遍历的结果来求第 K 个元素?该矩阵的特性很像 BST,因为该矩阵的扁平成一维数组的结果对于构建 BST 应该是比较理想的,不可能出现上图的这种情况。

我的思路:

  1. 将矩阵遍历,添加到一个新数组中;
  2. 将该序列构建 BST(新的序列对于构建 BST 是理想的);
  3. 对该 BST 中序遍历,结果放在一个新数组中;
  4. 得到该新数组的第 k 个元素。

先将该矩阵的元素来构建一颗二叉搜索树。以下是二叉树模型:

public class TreeNode {

    public int val;
    public TreeNode left;
    public TreeNode right;

    public TreeNode(int x) {
        val = x;
    }

    @Override
    public String toString() {
        return "TreeNode{" +
                "val=" + val +
                ", left=" + left +
                ", right=" + right +
                '}';
    }
}

通过深度优先遍历的方式,对矩阵的右上角向左向下遍历,为了不出现重复访问,设置一个 visited 访问标记。类似先序遍历:

public static void dfs2(int[][] matrix, int x, int y, boolean[][] visited, List<Integer> res) {
    if (x >= 0 && y >= 0 && x < matrix.length && y < matrix[0].length && !visited[x][y]) {
        visited[x][y] = true;
        res.add(matrix[x][y]);
        dfs2(matrix, x - 1, y, visited, res);
        dfs2(matrix, x, y + 1, visited, res);
    }
}

接下来就是构建 BST 的过程代码辣:

public static void buildBst(int n, TreeNode root) {
    if (n <= root.val) {
        if (root.left == null) {
            root.left = new TreeNode(n);
        } else {
            root = root.left;
            buildBst(n, root);
        }
    } else {
        if (root.right == null) {
            root.right = new TreeNode(n);
        } else {
            root = root.right;
            buildBst(n, root);
        }
    }

}

然后是对 BST 的中序遍历,BST 的中序遍历打印的结果,一定是升序的。

public static void dfs(TreeNode root, List<Integer> res) {
    if (root != null) {
        dfs(root.left, res);
        res.add(root.val);
        dfs(root.right, res);
    }
}

其中要注意 root 节点在构建之前就已经创建了,所以第一个遍历的数组下标从 1 开始遍历,完整的代码:

public class KthSmallestElementInaSortedMatrix {
    public static int kthSmallest(int[][] matrix, int k) {
        List<Integer> tmp = new ArrayList<>();
        List<Integer> res = new ArrayList<>();

        boolean[][] visited = new boolean[matrix.length][matrix[0].length];
        dfs2(matrix, matrix.length - 1, 0, visited, tmp);
        TreeNode root = new TreeNode(tmp.get(0));
        // 第一位是root节点,已经add了
        for (int i = 1; i < tmp.size(); i++) {
            buildBst(tmp.get(i), root);
        }
        dfs(root, res);
        return res.get(k - 1);
    }

    public static void buildBst(int n, TreeNode root) {
        if (n <= root.val) {
            if (root.left == null) {
                root.left = new TreeNode(n);
            } else {
                root = root.left;
                buildBst(n, root);
            }
        } else {
            if (root.right == null) {
                root.right = new TreeNode(n);
            } else {
                root = root.right;
                buildBst(n, root);
            }
        }


    }

    public static void dfs(TreeNode root, List<Integer> res) {
        if (root != null) {
            dfs(root.left, res);
            res.add(root.val);
            dfs(root.right, res);
        }
    }

    public static void dfs2(int[][] matrix, int x, int y, boolean[][] visited, List<Integer> res) {
        if (x >= 0 && y >= 0 && x < matrix.length && y < matrix[0].length && !visited[x][y]) {
            visited[x][y] = true;
            res.add(matrix[x][y]);
            dfs2(matrix, x - 1, y, visited, res);
            dfs2(matrix, x, y + 1, visited, res);
        }
    }

}
class TreeNode {

    public int val;
    public TreeNode left;
    public TreeNode right;

    public TreeNode(int x) {
        val = x;
    }

    @Override
    public String toString() {
        return "TreeNode{" +
                "val=" + val +
                ", left=" + left +
                ", right=" + right +
                '}';
    }
}

最后,当然性能拉胯了。。这只是一种奇葩做法。。

  • LeetCode

    LeetCode(力扣)是一个全球极客挚爱的高质量技术成长平台,想要学习和提升专业能力从这里开始,充足技术干货等你来啃,轻松拿下 Dream Offer!

    209 引用 • 72 回帖
  • Java

    Java 是一种可以撰写跨平台应用软件的面向对象的程序设计语言,是由 Sun Microsystems 公司于 1995 年 5 月推出的。Java 技术具有卓越的通用性、高效性、平台移植性和安全性。

    3165 引用 • 8206 回帖 • 1 关注
  • 二叉搜索树
    2 引用
  • 算法
    388 引用 • 254 回帖 • 22 关注

相关帖子

欢迎来到这里!

我们正在构建一个小众社区,大家在这里相互信任,以平等 • 自由 • 奔放的价值观进行分享交流。最终,希望大家能够找到与自己志同道合的伙伴,共同成长。

注册 关于
请输入回帖内容 ...