leetcode230：二叉搜索树中第K小的元素

题目链接

leetcode

题目描述

给定一个二叉搜索树的根节点 root 和一个整数 k ，设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

C++ 代码

#include <iostream>
#include <vector>
#include <stack>
#include <queue>
#include <unordered_map>
using namespace std;

// 二叉树结点的定义
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

const int NULL_NODE = -1;

// 辅助函数：创建二叉树
TreeNode* createTree(const vector<int>& nums) {
    if (nums.empty()) return nullptr;

    TreeNode* root = new TreeNode(nums[0]);
    queue<TreeNode*> queue;
    queue.push(root);
    int i = 1;
    while (!queue.empty() && i < nums.size()) {
        TreeNode* node = queue.front(); queue.pop();
        if (nums[i] != NULL_NODE) {
            node->left = new TreeNode(nums[i]);
            queue.push(node->left);
        }
        i ++;
        if (i < nums.size() && nums[i] != NULL_NODE) {
            node->right = new TreeNode(nums[i]);
            queue.push(node->right);
        }
        i ++;
    }

    return root;
}

// 辅助函数：层序遍历打印二叉树
void levelOrderTraversal(TreeNode* root) {
    if (root == nullptr) return;

    queue<TreeNode*> queue;
    queue.push(root);
    while (!queue.empty()) {
        TreeNode* node = queue.front(); queue.pop();
        cout << node->val << " ";
        if (node->left != nullptr) queue.push(node->left);
        if (node->right != nullptr) queue.push(node->right);
    }
    cout << endl;
}

// 辅助函数：释放二叉树
void deleteTree(TreeNode* root) {
    if (root == nullptr) return;
    deleteTree(root->left);
    deleteTree(root->right);
    delete root;
}

/*
    递归实现

    中序遍历整棵二叉搜索树，并用数组记录结点的值，输出第 k 个结点。
    
    时间复杂度：O(n)
        其中 n 为二叉树结点的数目。需要遍历整棵二叉搜索树。
    空间复杂度：O(n)
        数组需要 O(n) 的空间。
*/
class Solution_0 {
public:
    int kthSmallest(TreeNode* root, int k) {
        vector<int> res;
        dfs(root, res);
        return res[k - 1];
    }

    void dfs(TreeNode* node, vector<int>& res) {
        if (node == nullptr) return;

        dfs(node->left, res);
        res.push_back(node->val);
        dfs(node->right, res);
    }
};

/*
    递归实现，优化空间

    时间复杂度：O(n)
        其中 n 为二叉树结点的数目。需要遍历整棵二叉搜索树。
    空间复杂度：O(logn)
        在最坏的情况下，递归栈的深度可以达到树的高度。
        对于平衡二叉树，空间复杂度为 O(logn)。
*/
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        this->k = k;
        dfs(root);
        return res;
    }

    void dfs(TreeNode* node) {
        if (node == nullptr) return;

        dfs(node->left);
        if (-- k == 0) res = node->val;
        dfs(node->right);
    }

private:
    int res, k;
};

/*
    迭代实现

    在栈的辅助下，不用遍历整个树，找到答案后就停止。

    时间复杂度：O(n)
        其中 n 是二叉树的结点数量。
    空间复杂度：O(logn)
        其中 n 是二叉树的结点数量。
        树是一个平衡树时：O(logN)
        树是一个不平衡树时：O(N)
*/
class Solution_2 {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> stk;

        while (root || stk.size()) {
            while (root) {
                stk.push(root);
                root = root->left;
            }
            root = stk.top(); stk.pop();
            if (-- k == 0) return root->val;
            root = root->right;
        }

        return -1;
    }
};

int main() {
    // 示例输入
    Solution solution;
    vector<int> nums = {3, 1, 4, 0, 2};
    int k = 1;

    // 创建二叉搜索树
    TreeNode* root = createTree(nums);
    levelOrderTraversal(root);

    // 获取第 K 小的元素
    int result = solution.kthSmallest(root, k);
    // 打印结果
    cout << "Kth Smallest Element: " << result << endl;

    // 释放内存
    deleteTree(root);

    return 0;
}

Golang 代码

package main

import "fmt"

// 定义二叉树节点
type TreeNode struct {
    Val   int
    Left  *TreeNode
    Right *TreeNode
}

// 创建二叉树
func createTree(values []int) *TreeNode {
	if len(values) == 0 {
		return nil
	}
	root := &TreeNode{Val: values[0]}
	queue := []*TreeNode{root}
	i := 1
	for i < len(values) {
		node := queue[0]
		queue = queue[1:]

		if i < len(values) && values[i] != -1 {
			node.Left = &TreeNode{Val: values[i]}
			queue = append(queue, node.Left)
		}
		i++

		if i < len(values) && values[i] != -1 {
			node.Right = &TreeNode{Val: values[i]}
			queue = append(queue, node.Right)
		}
		i++
	}
	return root
}

// 打印二叉树（层序遍历）
func printTree(root *TreeNode) string {
	if root == nil {
		return "[]"
	}
	result := []int{}
	queue := []*TreeNode{root}
	for len(queue) > 0 {
		node := queue[0]
		queue = queue[1:]
		if node == nil {
			result = append(result, -1)
			continue
		}
		result = append(result, node.Val)
		queue = append(queue, node.Left)
		queue = append(queue, node.Right)
	}
	// 去掉尾部的 -1
	for len(result) > 0 && result[len(result)-1] == -1 {
		result = result[:len(result)-1]
	}
	return fmt.Sprintf("%v", result)
}

/*
    基本思路：递归实现
        二叉搜索树的中序遍历是一个升序序列。
        中序遍历，每访问一个节点就增加计数器，
        当计数器等于 k 时，当前节点的值就是第 k 小的元素。
    
    时间复杂度：O(n)
       其中 n 是树中节点的数量，需要遍历整个二叉搜索树。
    空间复杂度：O(logn)
        取决于栈的大小，它由树的高度决定。
        在最坏情况下，树是完全不平衡的，栈的大小为 O(n)。
        在平均情况下，树是完全平衡的，栈的大小为 O(logn)。
*/
func kthSmallest_0(root *TreeNode, k int) int {
    var res []int
    inorder(root, &res)
    return res[k - 1]
}

func inorder(node *TreeNode, res *[]int) {
    if node == nil {
        return
    }
    inorder(node.Left, res)
    *res = append(*res, node.Val)
    inorder(node.Right, res)
}

/*
    思路同上，优化空间
*/
var count int
var result int

func kthSmallest(root *TreeNode, k int) int {
    count = 0
    result = 0
    inorderTraversal(root, k)
    return result
}

func inorderTraversal(node *TreeNode, k int) {
    if node == nil {
        return
    }

    inorderTraversal(node.Left, k)
    count ++
    if count == k {
        result = node.Val
        return
    }
    inorderTraversal(node.Right, k)
}

/*
    基本思路：迭代实现
        使用一个栈来模拟中序遍历的过程。
        将当前节点的所有左子节点入栈，然后取出栈顶节点，对 k 进行递减操作并判断。
        然后处理当前节点的右子树。
        重复这个过程，找到答案后就停止。

    时间复杂度：O(n)
        其中 n 是二叉树的节点数量。
    空间复杂度：O(logn)
        其中 n 是二叉树的节点数量。
        树是一个平衡树时：O(logN)
        树是一个不平衡树时：O(N)
*/
func kthSmallest_2(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
}

func main() {
    // 测试用例
    testCases := []struct {
        values   []int
        k        int
        expected int
    }{
        {
            values:   []int{3,1,4,-1,2},
            k:        1,
            expected: 1,
        },
        {
            values:   []int{5,3,6,2,4,-1,-1,1},
            k:        3,
            expected: 3,
        },
    }

    for i, tc := range testCases {
        root := createTree(tc.values)
        fmt.Printf("Test Case %d, Input: values = %v, k = %d\n", i+1, printTree(root), tc.k)
        result := kthSmallest(root, tc.k)
        
        if result == tc.expected {
            fmt.Printf("Test Case %d, Output: %d, PASS\n", i+1, result)
        } else {
            fmt.Printf("Test Case %d, Output: %d, FAIL (Expected: %d)\n", i+1, result, tc.expected)
        }
    }
}