230 Kth Smallest Element in a BST

Given a binary search tree, write a functionkthSmallestto find thekth smallest element in it.
Note: You may assume k is always valid, 1 ? k ? BST's total elements.
Follow up: What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?
The Idea: Just do inorder traversal until the kth element is reached. Admitted, this is terrible python code (with the use of references), but it surprising scored in the top percentile range.
Complexity: O(N) time O(N) space
# Definition for a binary tree node.
# class TreeNode(object):
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class Solution(object):
def kthSmallest(self, root, k):
:type root: TreeNode
:type k: int
:rtype: int
count = [k]
stop = [False]
val = [0]
self.in_order_to_k(root, count, stop, val)
return val[0]
def in_order_to_k(self, root, cnt, stop, val):
:type root: TreeNode
:type cnt: int
:rtype: void
if (root and not stop[0]):
self.in_order_to_k(root.left, cnt, stop, val)
cnt[0] -= 1
if (cnt[0] == 0 and not stop[0] and root):
val[0] = root.val
stop[0] = True
self.in_order_to_k(root.right, cnt, stop, val)