Tuesday, January 20, 2015

Reverse a BST, find the kth largest element in BST

This is a FB interview question.
1. Reverse the BST;
2. find the kth largest element in O(n) time;
3. find the kth largest element in O(log n) time.

The first one is easy, recursively flip the left and right child of the tree.
The second one too, use in order traversal.
Now comes the third one. Here is an answer I found.
Here's just an outline of the idea:
Let each node in the BST have a field that returns the number of elements in its left and right subtree. Let the left subtree of node T contain only elements smaller than T and the right subtree only elements larger than or equal to T.
Now, suppose we are at node T:
  1. k == num_elements(left subtree of T), then the answer we're looking for is the value in nodeT
  2. k > num_elements(left subtree of T) then obviously we can ignore the left subtree, because those elements will also be smaller than the kth smallest. So, we reduce the problem to finding the k - num_elements(left subtree of T) smallest element of the right subtree.
  3. k < num_elements(left subtree of T), then the kth smallest is somewhere in the left subtree, so we reduce the problem to finding the kth smallest element in the left subtree.
This is O(log N) on average (assuming a balanced tree).

But the problem is, how to get the size of the tree? In the end, I wrote a Tree class. However, there is a bug in this class: I can only build the tree bottom up: otherwise the children of the subtree will not be added into the parent's children list.

import java.util.*;
public class Tree {
 private int size;
 TreeNode node;
 Tree left;
 Tree right;
 List children;
 public Tree() {
  children = new ArrayList();
  size = 0;
 }
 public Tree(int rootVal) {
  node = new TreeNode(rootVal);
  children = new ArrayList ();
  left = new Tree();
  right = new Tree();
  size++;
 }
 private void addChildren(Tree tree) {
  children.addAll(tree.children);
  size += tree.children.size();
 }
 public void buildLeft(Tree leftTree) {
  node.left = leftTree.node;
  left = leftTree;
  children.add(leftTree.node);
  size++;
  addChildren(leftTree);
 }
 
 public void buildRight(Tree rightTree) {
  node.right = rightTree.node;
  right = rightTree;
  children.add(rightTree.node);
  size++;
  children.addAll(rightTree.children);
  size += rightTree.children.size();
 }
 public int size() {
  return size;
 }

}
public class TreeNode {

  int val;
  TreeNode left;
  TreeNode right;
  //TreeNode parent;
  TreeNode(int x) { val = x;}

}


Here are the methods for the three problems:


public static void reverseBST (TreeNode root) {
  if (root == null || (root.left == null && root.right == null))
   return;
  TreeNode tmp = root.left;
  root.left = root.right;
  root.right = tmp;
  reverseBST(root.right);
  reverseBST(root.left);
 }
 static int count = 0;
 /**
  * O(n) complexity
  * @param root
  * @param k
  * @return
  */
 public static int kthLargest(Tree root, int k) {
  if (root == null)
   return -1;
  return findKth(root.node, k).val;
  
 }
 private static TreeNode findKth (TreeNode root, int k) {
  if (root == null)
   return root;
  TreeNode right = findKth(root.right, k);
  if (count == k)
   return right;
  ++count;
  if (count == k)
   return root;
  return findKth(root.left, k);
 }
 /**
  * O(log n) complexity
  * @param root
  * @param k
  * @return
  */
    public static TreeNode findKth(Tree root, int k) {
     if (root.right.size() == k - 1)
      return root.node;
     else if (root.right.size() < k - 1)
      return findKth(root.left, k - root.right.size() - 1);
     else
      return findKth(root.right, k);
  
 }

No comments:

Post a Comment