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

2018-10-20  本文已影响23人  one_zheng

给定一个二叉搜索树,编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明:
你可以假设 k 总是有效的,1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

image.png

示例 2:

image.png

进阶:
如果二叉搜索树经常被修改(插入/删除操作)并且你需要频繁地查找第 k 小的值,你将如何优化 kthSmallest 函数?

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func kthSmallest(root *TreeNode, k int) int {
    if root == nil {
        return 0
    }
    array := make([]int, 0)
    stack := NewStack()
    stack.Push(root)
    for {
        if stack.Empty() {
            break
        }
        root = stack.Pop().(*TreeNode)
        array = append(array, root.Val)
        if root.Right != nil {
            stack.Push(root.Right)
        }
        if root.Left != nil {
            stack.Push(root.Left)
        }
    }
    sort.Ints(array)

    if k < 0 || k-1 >= len(array) {
        return 0
    }
    return array[k-1]
}

type Stack struct {
    list *list.List
}

func NewStack() *Stack {
    list := list.New()
    return &Stack{list}
}

func (stack *Stack) Push(value interface{}) {
    stack.list.PushBack(value)
}

func (stack *Stack) Pop() interface{} {
    e := stack.list.Back()
    if e != nil {
        stack.list.Remove(e)
        return e.Value
    }
    return nil
}

func (stack *Stack) Peak() interface{} {
    e := stack.list.Back()
    if e != nil {
        return e.Value
    }

    return nil
}

func (stack *Stack) Len() int {
    return stack.list.Len()
}

func (stack *Stack) Empty() bool {
    return stack.list.Len() == 0
}


上一篇下一篇

猜你喜欢

热点阅读