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

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 +
                '}';
    }
}
​

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