How to select a binary tree node uniformly at random

The exercise I'm trying to solve is


You are implementing a binary search tree class from scratch, which, in addition, to insert, find and delete, has a method getRandomNode() which returns a random node from the tree. All nodes should be equally likely to be chosen. Design and implement an algorithm for getRandomNode(), and explain how you would implement the rest of the methods.

The answer from the book is:

1  class TreeNode {
 2    private int data;
 3    public TreeNode left;
 4    public TreeNode right;
 5    private int size = 0;
 ...
12    public TreeNode getRandomNode() {
13      int leftSize = left == null ? 0 : left.size();
14      Random random = new Random();
15      int index = random.nextInt(size);
16      if (index < leftSize) {
17        return left.getRandomNode();
18      } else if (index == leftSize) {
19        return this;
20      } else {
21        return right.getRandomNode();
22      }
23    }
...
55  }

But here is the problem. With this algorithm, I don't see how nodes are equally likely to be chosen. In line 16, it says if leftsize > index, where index is a number from 0 to size, then the algorithm will continue with the left node, otherwise the right node. It only works when the tree has a depth of 2. When the tree is taller, the probability of each node being chosen will not be equal.

Am I wrong? Does this algorithm work?

The algorithm works just fine.

Note that each node's size field tells you the total number of nodes in the subtree rooted at that node. Throughout this answer, I'm just going to assume that a child's size is zero if the child is null. Thus, for any node in the tree, we have

this.size == left.size + 1 + right.size

Consider the root of the tree. If we choose a node of the tree uniformly at random, then it should be in the left subtree with probability left.size/size, it should be the root node with probability 1/size and it should be in the right subtree with probability right.size/size. By the equation above, these three probabilities add up to 1. If you look at lines 15–22, the algorithm implements exactly this probability distribution. It generates a uniform random integer in the range 0 .. size-1. If this number is in the range 0 .. left.size-1, the algorithm picks a random node from the left subtree; if it's equal to left.size, it picks the root and, otherwise, it picks a random node from the right subtree.

And the argument I gave doesn't depend on the node you consider being the root, so the recursive calls behave as they should. The algorithm doesn't care about the depth of the nodes: it just recurses down the subtree and chooses subtrees with the correct probability.

StackExchange Computer Science Q#83540, answer score: 7