Solution: Drawing a Ball at Random

A bag promises nothing about order, so there is no need to shift. We take the last ball and drop it into the hole. This is the swap-with-the-last removal we saw when we looked at what remove promises. For DynamicArray we had to rule it out, because it reorders the elements and the contract forbids that.

rng.nextInt(size) gives us a value in [0, size), and every value in that range is equally likely. That is what makes the draw uniform over the balls currently in the bag. Note that we pass size and not balls.length, because the slots past size do not hold balls.

public BallColor remove() {
  if (size == 0) {
    throw new IllegalStateException("the bag is empty");
  }
  int i = rng.nextInt(size);    // pick a slot uniformly at random in [0, size)
  BallColor drawn = balls[i];   // this is the ball we hand back
  balls[i] = balls[size - 1];   // fill the gap with the last ball
  size--;
  balls[size] = null;           // clear the now-unused slot
  return drawn;
}

Here is a trace of one draw. The bag starts with four balls, and the random pick lands on slot 0, which holds a BLACK ball. That is the ball we will return:

┌───────┬────────┬───────┬────────┐
│ BLACK │  BLUE  │ BLACK │  BLUE  │
└───────┴────────┴───────┴────────┘
    ↑
  draw

We read that BLACK ball out, which leaves a hole at slot 0. To close it, we move the last ball, the BLUE at slot 3, into the hole:

┌───────┬────────┬───────┬────────┐
│  BLUE │  BLUE  │ BLACK │  BLUE  │
└───────┴────────┴───────┴────────┘
    ↑
 filled

Now slot 0 holds the ball that used to be at the end, and slot 3 still refers to that same ball. We drop size to 3 and clear that freed slot so it no longer refers to a ball:

┌───────┬────────┬───────┬───────┐
│  BLUE │  BLUE  │ BLACK │   _   │
└───────┴────────┴───────┴───────┘
                             ↑
                          cleared

Then we return the drawn BLACK ball.

The whole operation is one read, one copy, one decrement, and one clear. That is the same amount of work no matter which slot the draw lands on. Compare that to the dynamic array, where removing an early element means shifting every later element down. The bag is cheaper because it does not promise to keep the balls in any order. There is no order for the code to maintain.