230. 二叉搜索树中第 K 小的元素

题目描述

输入一棵二叉搜索树的根节点 root 和一个整数 k。树用层序数组表示，空位记为 null。

把树中所有节点值从小到大排列，返回排在第 k 位的值。k 从 1 开始数。

示例 1

输入：root = [3,1,4,null,2], k = 1
输出：1

解释：升序排列为 1,2,3,4，第 1 个是 1。

示例 2

输入：root = [5,3,6,2,4,null,null,1], k = 3
输出：3

解释：升序排列为 1,2,3,4,5,6，第 3 个是 3。

数据范围

树的节点总数记为 n。
1 <= k <= n <= 10^4
0 <= Node.val <= 10^4

进阶：如果这棵树会被频繁插入和删除节点，同时又要频繁查询第 k 小的值，应该怎样改进？

解题思路：迭代中序遍历

执行迭代中序遍历，每弹出一个节点就将 k 减一，减到零时返回。

为什么正确

中序已访问节点严格按升序排列，k 表示还需访问多少个节点。

复杂度分析

最坏时间 O(n)，提前停止时 O(h+k)，空间 O(h)。

易错点

k 从 1 开始计数；在弹栈访问时减一，而不是入栈时。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def kthSmallest(self, root, k):
        stack = []
        cur = root
        while cur or stack:
            while cur:
                stack.append(cur)
                cur = cur.left
            cur = stack.pop()
            k -= 1
            if k == 0:
                return cur.val
            cur = cur.right

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode cur = root;
        while (cur != null || !stack.isEmpty()) {
            while (cur != null) {
                stack.push(cur);
                cur = cur.left;
            }
            cur = stack.pop();
            k--;
            if (k == 0)
                return cur.val;
            cur = cur.right;
        }
        return -1;
    }
}

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> stk;
        TreeNode* cur = root;
        while (cur || !stk.empty()) {
            while (cur) {
                stk.push(cur);
                cur = cur->left;
            }
            cur = stk.top(); stk.pop();
            k--;
            if (k == 0)
                return cur->val;
            cur = cur->right;
        }
        return -1;
    }
};

func kthSmallest(root *TreeNode, k int) int {
	stack := []*TreeNode{}
	cur := root
	for cur != nil || len(stack) > 0 {
		for cur != nil {
			stack = append(stack, cur)
			cur = cur.Left
		}
		cur = stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		k--
		if k == 0 {
			return cur.Val
		}
		cur = cur.Right
	}
	return -1
}

var kthSmallest = function(root, k) {
    const stack = [];
    let cur = root;
    while (cur || stack.length) {
        while (cur) {
            stack.push(cur);
            cur = cur.left;
        }
        cur = stack.pop();
        k--;
        if (k === 0)
            return cur.val;
        cur = cur.right;
    }
};

function kthSmallest(root: TreeNode | null, k: number): number {
    const stack: TreeNode[] = [];
    let cur = root;
    while (cur || stack.length) {
        while (cur) {
            stack.push(cur);
            cur = cur.left;
        }
        cur = stack.pop()!;
        k--;
        if (k === 0)
            return cur.val;
        cur = cur.right;
    }
    return -1;
}

载入编辑器，自己跑一遍 →

其他解法

迭代中序遍历可以提前停止，最常用；子树节点计数把预处理和查询分开，适合多次查询；Morris 遍历用临时线索换来 O(1) 空间。

解法 2：子树节点计数时间 O(n)，空间 O(n)

先递归统计以每个节点为根的子树节点数，存入哈希表。从根出发，设左子树有 left 个节点：k ≤ left 时去左子树；k = left + 1 时当前节点就是答案；否则令 k 减去 left + 1 后去右子树。

单次查询要先花 O(n) 统计。统计结果可以复用，之后每次查询只需 O(h)，适合频繁查询第 k 小的进阶场景。

class Solution:
    def kthSmallest(self, root, k):
        size = {}

        def count(node):
            if not node:
                return 0
            size[node] = 1 + count(node.left) + count(node.right)
            return size[node]

        count(root)
        node = root
        while node:
            left = size.get(node.left, 0) if node.left else 0
            if k <= left:
                node = node.left
            elif k == left + 1:
                return node.val
            else:
                k -= left + 1
                node = node.right
        return -1

class Solution {
    private final Map<TreeNode, Integer> size = new HashMap<>();

    private int count(TreeNode node) {
        if (node == null) return 0;
        int s = 1 + count(node.left) + count(node.right);
        size.put(node, s);
        return s;
    }

    public int kthSmallest(TreeNode root, int k) {
        count(root);
        TreeNode node = root;
        while (node != null) {
            int left = node.left != null ? size.get(node.left) : 0;
            if (k <= left) {
                node = node.left;
            } else if (k == left + 1) {
                return node.val;
            } else {
                k -= left + 1;
                node = node.right;
            }
        }
        return -1;
    }
}

class Solution {
    unordered_map<TreeNode*, int> size;

    int count(TreeNode* node) {
        if (!node) return 0;
        return size[node] = 1 + count(node->left) + count(node->right);
    }

public:
    int kthSmallest(TreeNode* root, int k) {
        count(root);
        TreeNode* node = root;
        while (node) {
            int left = node->left ? size[node->left] : 0;
            if (k <= left) {
                node = node->left;
            } else if (k == left + 1) {
                return node->val;
            } else {
                k -= left + 1;
                node = node->right;
            }
        }
        return -1;
    }
};

func kthSmallest(root *TreeNode, k int) int {
	size := map[*TreeNode]int{}
	var count func(node *TreeNode) int
	count = func(node *TreeNode) int {
		if node == nil {
			return 0
		}
		size[node] = 1 + count(node.Left) + count(node.Right)
		return size[node]
	}
	count(root)
	node := root
	for node != nil {
		left := size[node.Left]
		if k <= left {
			node = node.Left
		} else if k == left+1 {
			return node.Val
		} else {
			k -= left + 1
			node = node.Right
		}
	}
	return -1
}

var kthSmallest = function(root, k) {
    const size = new Map();
    const count = (node) => {
        if (!node) return 0;
        const s = 1 + count(node.left) + count(node.right);
        size.set(node, s);
        return s;
    };
    count(root);
    let node = root;
    while (node) {
        const left = node.left ? size.get(node.left) : 0;
        if (k <= left) {
            node = node.left;
        } else if (k === left + 1) {
            return node.val;
        } else {
            k -= left + 1;
            node = node.right;
        }
    }
    return -1;
};

function kthSmallest(root: TreeNode | null, k: number): number {
    const size = new Map<TreeNode, number>();
    const count = (node: TreeNode | null): number => {
        if (!node) return 0;
        const s = 1 + count(node.left) + count(node.right);
        size.set(node, s);
        return s;
    };
    count(root);
    let node = root;
    while (node) {
        const left = node.left ? size.get(node.left)! : 0;
        if (k <= left) {
            node = node.left;
        } else if (k === left + 1) {
            return node.val;
        } else {
            k -= left + 1;
            node = node.right;
        }
    }
    return -1;
}

解法 3：Morris 中序遍历时间 O(n)，空间 O(1)

当前节点有左子树时，找到左子树的最右节点。最右节点的右指针为空就让它指向当前节点，然后进入左子树；不为空说明左子树已访问完，恢复指针、访问当前节点、进入右子树。没有左子树时直接访问当前节点并进入右子树。每访问一个节点 k 减一，减到零时记录答案。

不用栈，额外空间 O(1)。遍历中会临时修改树指针，为了恢复树结构要走完整棵树，代码也更难写对。

class Solution:
    def kthSmallest(self, root, k):
        ans = -1
        cur = root
        while cur:
            if cur.left:
                pre = cur.left
                while pre.right and pre.right is not cur:
                    pre = pre.right
                if pre.right is None:
                    pre.right = cur
                    cur = cur.left
                    continue
                pre.right = None
            k -= 1
            if k == 0:
                ans = cur.val
            cur = cur.right
        return ans

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        int ans = -1;
        TreeNode cur = root;
        while (cur != null) {
            if (cur.left != null) {
                TreeNode pre = cur.left;
                while (pre.right != null && pre.right != cur) pre = pre.right;
                if (pre.right == null) {
                    pre.right = cur;
                    cur = cur.left;
                    continue;
                }
                pre.right = null;
            }
            if (--k == 0) ans = cur.val;
            cur = cur.right;
        }
        return ans;
    }
}

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int ans = -1;
        TreeNode* cur = root;
        while (cur) {
            if (cur->left) {
                TreeNode* pre = cur->left;
                while (pre->right && pre->right != cur) pre = pre->right;
                if (!pre->right) {
                    pre->right = cur;
                    cur = cur->left;
                    continue;
                }
                pre->right = nullptr;
            }
            if (--k == 0) ans = cur->val;
            cur = cur->right;
        }
        return ans;
    }
};

func kthSmallest(root *TreeNode, k int) int {
	ans := -1
	cur := root
	for cur != nil {
		if cur.Left != nil {
			pre := cur.Left
			for pre.Right != nil && pre.Right != cur {
				pre = pre.Right
			}
			if pre.Right == nil {
				pre.Right = cur
				cur = cur.Left
				continue
			}
			pre.Right = nil
		}
		k--
		if k == 0 {
			ans = cur.Val
		}
		cur = cur.Right
	}
	return ans
}

var kthSmallest = function(root, k) {
    let ans = -1;
    let cur = root;
    while (cur) {
        if (cur.left) {
            let pre = cur.left;
            while (pre.right && pre.right !== cur) pre = pre.right;
            if (!pre.right) {
                pre.right = cur;
                cur = cur.left;
                continue;
            }
            pre.right = null;
        }
        if (--k === 0) ans = cur.val;
        cur = cur.right;
    }
    return ans;
};

function kthSmallest(root: TreeNode | null, k: number): number {
    let ans = -1;
    let cur = root;
    while (cur) {
        if (cur.left) {
            let pre: TreeNode = cur.left;
            while (pre.right && pre.right !== cur) pre = pre.right;
            if (!pre.right) {
                pre.right = cur;
                cur = cur.left;
                continue;
            }
            pre.right = null;
        }
        if (--k === 0) ans = cur.val;
        cur = cur.right;
    }
    return ans;
}

二叉树专题的其他题目

查看二叉树全部题目 →