Tuesday, August 18, 2015

Reservoir Sampling - Tree

Problem

Get n random nodes from a binary tree, such that each node will be returned with equal probability.

Solution

import java.util.ArrayList;
import java.util.Random;

public class Tree {
 public int node;
 public Tree left;
 public Tree right;
}
public class NodesAtRandom {
 
 int index = 0;
 ArrayList getRandomNodes(Tree t, int n){
  ArrayList ret = new ArrayList();
  index =0;
  preOrder(ret, t, n);
  return ret;
 }
 
 private void preOrder(ArrayList ret, Tree t, int n) {
  if(t==null) return;
  index++;
  if(ret.size() < n) ret.add(t.node);
  else {
   Random r = new Random();
   int rVal = r.nextInt(index);
   if(rVal < n) ret.set(rVal, t.node);
  }
  preOrder(ret, t.left, n);
  preOrder(ret, t.right, n);
 }

 public static void main(String[] args){
  Tree t = TreeTestData.getSampleTree1();
  TreeOperations.printTree(t);
  NodesAtRandom ob = new NodesAtRandom();
  ArrayList output = ob.getRandomNodes(t, 2);
  System.out.println(output);
 }
}

No comments:

Post a Comment