Why subtree size is used in algorithm to select random node in a binary tree?

Viewed 125

I stumbled upon several implementations of algorithm to select a random node from a binary tree, and they all use subtree size property. However, I don't understand why knowing subtree size helps. Here's implementation A and B.

The general idea can be described in this pseudocode:

Node getRandomNode() {
    Random seed = new Random;
    int random = random.nextInt(this.size);

    if (this.left.size == random)
        return this;
  
    if (this.left.size > random) {
        return this.left.getRandomNode();
    } else {
        return this.right.getRandomNode();
    }
}

I assume that the getRandomNode() method is always called on the root of the tree.

As far as I see the randomality of choosing a node depends solely on the randomality of random integer. In certain cases if random repeats itself then getRandomNode() will return the same node.

Also I don't understand how the algorithm can work in a complete binary tree. The size of each node subtree will always be an even number. So if random is an even number the algorithm will not have a match on this line:

if (this.left.size == random)
        return this;

and will only produce a random number if random is an even number, which means that the algorithm is not random.

Am I missing something?

If I had to implement the algorithm I would just store an additional linked list where each node would have an index (like ArrayList in Java) and then return the node from the list whose index equals a random number returned from random.nextInt().

4 Answers

Thanks for outlining your thought process here. Let's go through what you're saying here one step at a time.

I assume that the getRandomNode() method is always called on the root of the tree.

It is always called on the root of some subtree, but it's not necessarily going to be the overall tree root. For example, consider this simple tree:

                  A
                 / \
                B   C
                   / \
                  D   E

Here, the first call will be to A.getRandomNode(), which will pick a random node out of the subtree rooted at A. After that, there's a 1/5 chance that we ask for a random node out of the left subtree, a 3/5 chance that we ask for a random node out of the right subtree, and a 1/5 chance that we stop and return A. Let's say, for the sake of argument, that the random number generated here is 4 and that we decide to look in the right subtree. That means that the next call is to C.getRandomNode(), which asks to pick a random node out of this tree:

             C
            / \
           D   E

Here, we'll generate a random number such that there's a 1/3 chance of asking for a random node from the left subtree, a 1/3 chance of asking for a random node from the right subtree, and a 1/3 chance of stopping with C. For the sake of argument, suppose we generate the random number 0 and go to the left. That calls D.getRandomNode(), which asks for a random node out of this subtree:

               D

This call will always return D, since that's the only option.

Hopefully this gives you a bit more of a sense of how the algorithm works. At each point in time, we want to have an equal chance of choosing any element. Therefore, we decide whether to stop, go left, or go right, weighting our choice based on the number of nodes in each of the option.

As far as I see the randomality of choosing a node depends solely on the randomality of random integer. In certain cases if random repeats itself then getRandomNode() will return the same node.

You're correct that the only source of randomness here is which random integer gets chosen. And that would make sense, since the tree itself isn't random.

That being said, the implementation that you've shown here makes a separate random calculation at each step in walking down the tree. Therefore, multiple different random choices all have to go the same way in order to get the same node as the reuslt.

Also I don't understand how the algorithm can work in a complete binary tree. The size of each node subtree will always be an even number. So if random is an even number the algorithm will not have a match on this line:

if (this.left.size == random)
     return this;

and will only produce a random number if random is an even number, which means that the algorithm is not random.

You're correct that if we just generated a single random number and used that number throughout the process, then yes, you'd run into trouble with something like this. However, that's not how the algorithm works. Instead, each time we make a recursive call to getRandomNode(), we pick a new random number and use that to do our routing. Therefore, regardless of whether the first number we picked is even or odd, we'll eventually get to a spot where the only option is to return the current node, at which point we return something.

If I had to implement the algorithm I would just store an additional linked list where each node would have an index (like ArrayList in Java) and then return the node from the list whose index equals a random number returned from random.nextInt().

You absolutely could solve the problem this way. If the tree you're working with never changes - nothing is ever added or removed - then this is a perfectly reasonable strategy.

The reason that the approach of tagging subtree sizes is useful is that it lets you sample from the tree randomly even if nodes are being added or removed. Specifically, you can fairly quickly adjust the numbers stored in each node in response to an insertion or removal of a node, much faster than having to rebuild or shuffle elements around in an ArrayList were the tree to change shape. Check out the order statistic tree for more details about how to do this.

You could convert the tree to an array/linked list if you wanted to. Generate a random number and get the element stored at that index. But converting the tree would also take an extra O(N) step.
Now, coming to the subtree size approach. This algorithm is also doing the same thing recursively: mapping a value to each node value (analogous to the index in array). The mapping basically becomes

0 -> root node
1 - left_subtree_size -> for all left subtree nodes
left_subtree_size+1 - right_subtree_size -> for all right subtree nodes

Since each node gets mapped to only one value, there's an equal chance of each node being selected (obviously, assuming all values have equal probability under the random() function).

That algorithm, after some minor cleanup, is a clever and elegant use of conditional probability. Conditional probability says that if some event A is contingent/conditional on the occurrence of some other event B, then P{A} = P{A | B} * P{B} — in words, the probability of A is equal to the probability of A given that B occurred times the probability of B occurring. If B is itself contingent on C, this can cascade into P{A} = P{A | B} * P{B} = P{A | B} * P{B | C} * P{C}, and so on. You end up multiplying all of the probabilities of all the conditions that got you to the final state.

The algorithm in question gives us the conditional probabilities at a single stage, but it's easiest to see how it selects random tree nodes with equal probabilities by tracking it through a concrete example. I have a concrete implementation presented below. Since you didn't specify a language tag, I chose to use Ruby. (Note: this is not written in such a way to make Ruby purist's hearts go pit-a-pat, it is written to be readable by non-Rubyists, be self-contained and executable, and to illustrate how the conditional probabilities work.)

First, let's introduce a binary search tree node that tracks subtree sizes:

class BT_node
  # Make the following instance vars read-only accessible by name
  attr_reader :value, :left, :right, :size

  # Instantiate a new BT_node to hold the assigned value
  def initialize(value)
    @value = value
    # new values become leaves (see below), so both children are nil.
    @left = nil
    @right = nil
    # Assumption: size is the number of BT_node's in the subtree rooted
    # at, including, the current BT_node.  Starts at 1 since all newly
    # added nodes are leaves.
    @size = 1
  end

  def add_value(new_value)
    # Note: not checking, but this assumes all entries are unique.  In other
    # words, at each point a new_value belongs in either the left or right
    # subtree of the current node, it's never equal to an existing node's
    # value, so all new_value's become leaves in the tree when inserted.
    if new_value < @value
      if @left
        @left.add_value(new_value)
      else
        @left = BT_node.new(new_value)
      end
    else
      if @right
        @right.add_value(new_value)
      else
        @right = BT_node.new(new_value)
      end
    end
    @size += 1
  end

  # In-order traversal and printing of tree
  def treeprint
    @left.treeprint if @left
    puts "#{@value}, #{@size}"
    @right.treeprint if @right
  end

  # Working implementation of the random node picker.  Note how I carefully
  # prevented spell-check from changing that to nose.
  def random_pick
    rnd_value = rand(@size)
    lsize = @left ? @left.size : 0  # left.size if left exists, 0 otherwise
    return @value if rnd_value == lsize     # probability is 1/@size
    if rnd_value < lsize                    # probability is lsize/(@size-1)
      return @left.random_pick
    else                            # probability is right.size/(size-1)
      return @right.random_pick
    end
  end
end

I added comments to clarify assumptions I was making. I hope the setup is fairly obvious if you understand binary search trees, so let's focus on the random_pick. Given that we are at a particular node in the tree, by construction @size will tell you how many nodes there are in the current subtree (including itself). We then generate a number uniformly in the range [0,...,@size-1], inclusive, so each value can occur with probability 1/@size. Now for the slightly weird but mathematically correct part. (If this part bothers you, see the alternative implementation at the end.) We compare the generated rnd_value to the size of the left sub-tree, if it exists, or zero if there is no left child. Either way, there's a 1/@size chance of a match, in which case we return the current node's value. If we didn't match, we well end up choosing one of the @size - 1 remaining values in this subtree. There are lsize values in the inclusive range [0,...,lsize-1], so our destiny lies in the left subtree with probability lsize/(@size-1), or in the right subtree with the remaining probability (@size-1-lsize)/(@size-1).

That's the snapshot at a single level of the recursion, but let's get the full picture using conditional probability. That's easiest to track with a concrete example. Consider the specific binary tree constructed using following sequence of commands:

root = BT_node.new(3)
root.add_value(1)
root.add_value(2)
root.add_value(7)
root.add_value(5)
root.add_value(4)
root.add_value(6)
root.add_value(8)

That sequence of add_value's produces the following tree, where entries are of the form "value:size":

        3:8
      /     \
   1:2       7:5
     \      /   \
     2:1  5:3   8:1
         /   \
       4:1   6:1

Let's start simple, and ask how random_pick would yield a 3. We start at the root, and generate one of the 8 values value in the range [0,...,7]. If it happens to be 2, the size of the root's left subtree, we're done. This happens with probability 1/8. If you had not picked the root there are 7 nodes remaining for consideration, 2 of which are in the left subtree and 5 of which are in the right. You pick the left with probability 2/7 and the right with probability 5/7 to give all remaining nodes an equal chance.

Now let's look at a tougher one where we'll have to invoke conditional probability. How would random_pick yield a 5? To pick a 5, we would start at the root, fail to pick it, go to its right subtree, fail to pick the 7, go to its left subtree, and pick the 5. The probabilities at each stage are as follows:

  • fail to pick 3: 7/8
  • go to right sub-tree: 5/7
  • fail to pick 7: 4/5
  • go to left sub-tree: 3/4
  • pick the 5: 1/3

Conditional probability tells us that the probability of getting a 5 is product of all the probabilities along this contingency path: (7/8)*(5/7)*(4/5)*(3/4)*(1/3) = 1/8. Notice how the numerator of each term is cancelled by the denominator of the next.

Pick any other node in the tree, and following the same sort of conditional logic you will find that the probabilities are all 1/8. As you can see, the tree does not have to be balanced, the conditional probabilities guarantee that each node in the tree has an equal chance of being the one selected.

Here's the rest of the program, which confirms that the sampling is uniform across a million trials.

# # Confirmation of tree - uncomment to inspect tree structure
# p root
# root.treeprint

# # For analysis in a stats package
# 1_000_000.times { puts root.random_pick }

# For analysis by eye - Expected count for each value is 125000
h = Hash.new(0)  # default value set to zero for all keys
# Accumulate how many times we pick each value over 1 million picks
1_000_000.times { h[root.random_pick] += 1 }
p h  # inspect the hash

# Produces, e.g.:
#
# {4=>124824, 2=>125312, 8=>125298, 5=>124593, 6=>124562, 1=>125272, 7=>124857, 3=>125282}
#

I calculated the chi-squared statistic for this sample under the null hypothesis that the outcomes were equally likely. The expected value of a chi-square with 7 degrees of freedom is 7, and large deviations from expected will increase the statistic and lead to rejection of the null hypothesis. Since my test statistic was 5.9886, we solidly fail to reject the null hypothesis that the algorithm has an equal likelihood of generating each of the values in the tree.


If it really bothers you to compare the rnd_value to the left sub-tree's size, you can pick the current node when rnd_value == 0. You then have to adjust the comparison to less than or equal rather than less than to get the correct probability of selecting the left subtree. Here it is if you're interested:

# Alternate version if you don't like basing choice of self on left.size
def random_pick
  rnd_value = rand(@size)
  return @value if rnd_value == 0
  lsize = @left ? @left.size : 0
  if rnd_value <= lsize
    return @left.random_pick
  else
    return @right.random_pick
  end
end

Lastly, if you want the actual node rather than the value stored, change return @value to return self in either version of random_pick.

Adapted reservoir sampling (Knuths "Algorithm R") to tree search. You don't need to know the tree size or the sizes of the subtrees in advance. (but, unfortunately, you'll have to visit every node ...)


#include <stdio.h>
#include <stdlib.h>
#include <time.h>

struct node{
        struct node *left, *right;
        unsigned value;
        };

unsigned urand(unsigned lim)
{
unsigned rnd;

rnd = random();
rnd %= lim; // Biased!

return rnd;
}

struct node *tree_sample(struct node *tree, struct node *samp, unsigned *counter)
{
unsigned rnd;

if (!tree) return samp;

*counter += 1;
if (*counter== 1) { samp = tree; }
else    {
        rnd = urand(*counter);
        if (rnd==0) {samp=tree;}
        }

samp = tree_sample(tree->left, samp, counter);
samp = tree_sample(tree->right, samp, counter);

return samp;
}

struct node *node_new(unsigned value)
{
struct node *this;

this = malloc(sizeof *this);
this->left = NULL;
this->right = NULL;
this->value = value;

return this;
}

struct node **tree_find(struct node **tree, unsigned value)
{
while (*tree) {
        int cmp;
        cmp = value - (*tree)->value;
        if (!cmp) break;
        tree = (cmp < 0) ? &(*tree)->left : &(*tree)->right;
        }
return tree;
}

struct node *root = NULL;

int main(void)
{
unsigned count;
struct node *this;

        // Create some noise ...
for (count = 0; count < 10; count++) {
        struct node **pp;
        this = node_new(random());
        pp = tree_find(&root, this->value);
        *pp = this;
        }

srand(time(NULL));
count = 0;

        // pick a sample
this = tree_sample (root, NULL, &count);
printf("Count=%u\n", count);
if (this) printf("Val=%u\n", this->value);

return 0;
}
                                              
Related