{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "

230. Kth Smallest Element in a BST

\n", "
\n", "\n", "\n", "

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

\n", "\n", "

 

\n", "\n", "

Example 1:

\n", "\n", "
Input: root = [3,1,4,null,2], k = 1\n",
    "   3\n",
    "  / \\\n",
    " 1   4\n",
    "  \\\n",
    "   2\n",
    "Output: 1
\n", "\n", "

 

\n", "

Example 2:

\n", "\n", "
Input: root = [5,3,6,2,4,null,null,1], k = 3\n",
    "       5\n",
    "      / \\\n",
    "     3   6\n",
    "    / \\\n",
    "   2   4\n",
    "  /\n",
    " 1\n",
    "Output: 3\n",
    "
\n", "\n", "

 

\n", "

Constraints:

\n", "\n", "\n", "\n", "\n", "\n", "

 

\n", "Source \n", "
" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "# Definition for a binary tree node.\n", "class TreeNode:\n", " def __init__(self, val=0, left=None, right=None):\n", " self.val = val\n", " self.left = left\n", " self.right = right" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Code

" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "def kth_smallest(root, k):\n", " \"\"\"Recursive approach\n", " Time complexity: O(H+k), where H is a tree height. The stack contains at least H+k elements,\n", " since before starting to pop out one has to go down to a leaf.\n", " O(log⁡N+k) for the balanced tree\n", " O(N+k) for completely unbalanced tree\n", " Space complexity: O(H) to keep the stack, where H is a tree height.\n", " O(N) in the worst case of the skewed tree\n", " O(log⁡N) in the average case of the balanced tree\n", "\"\"\"\n", " def inorder(node):\n", " if node is None:\n", " return\n", " nonlocal k, traversal\n", " if len(traversal) >= k:\n", " return\n", " inorder(node.left)\n", " traversal.append(node.val)\n", " inorder(node.right)\n", "\n", " traversal = []\n", " inorder(root)\n", " return traversal[k-1] # !!! traversal[-1] don't work\n", " # # (always possible to add up to 2 more node after k)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "

Follow up #1:

\n", "

Solve it both recursively and iteratively

\n", "

Code

" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def kth_smallest(root, k):\n", " \"\"\"iterative approach\"\"\"\n", " stack = []\n", " traversal = []\n", " node = root\n", " while node or stack:\n", " while node: # go to the leftmost child\n", " if len(traversal) == k:\n", " break\n", " stack.append(node)\n", " node = node.left\n", " # if no more left child, get the 1st right node and check for leftmost child again\n", " \n", " if len(traversal) == k:\n", " break\n", " node = stack.pop()\n", " traversal.append(node.val)\n", " node = node.right\n", "\n", " return traversal[-1]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def kth_smallest(root, k):\n", " \"\"\"alternative iterative solution\"\"\"\n", " if not root:\n", " return None\n", "\n", " stack = []\n", " traversal = []\n", " node = root\n", " while len(traversal) < k:\n", " if node:\n", " stack.append(node)\n", " node = node.left\n", " elif stack and not node:\n", " node = stack.pop()\n", " traversal.append(node.val)\n", " node = node.right\n", " else:\n", " break\n", "\n", " return traversal[-1]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from itertools import islice\n", "\n", "def kth_smallest(root, k):\n", " \"\"\"alternative iterative solution using a generator\"\"\"\n", " def inorder(node):\n", " if not node:\n", " return\n", "\n", " yield from inorder(node.left)\n", " yield node.val\n", " yield from inorder(node.right)\n", "\n", " \n", " return next(islice(inorder(root), k-1, k))\n", " # return list(islice(inorder(root), k-1, k))[0] # alternatively" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "

Follow up #2:

\n", "

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?

\n", "

Code

" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# In a BST Insert and delete have a time complexity of O(H), where H = height (= log N for balanced tree)\n", "\n", "# Without any optimisation insert/delete + search of kth element has O(2H+k)\n", "\n", "# We could combine the BST with a double linked list. Then it would have:\n", "# O(H) time for the insert and delete.\n", "# O(k) for the search of kth smallest.\n", "\n", "# The overall time complexity for insert/delete + search of kth smallest is O(H+k) instead of O(2H+k)\n", "\n", "# Time complexity: O(H+k)\n", "# O(log⁡N+k)in average case\n", "# O(N+k) in worst case\n", "# Space complexity: O(N) to keep the linked list." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.2" } }, "nbformat": 4, "nbformat_minor": 1 }