## Ball Pit

We need to build a data structure that simulates a ball pit. It needs to support 3 operations:

`void Insert(int color)`

- Add a ball of color \(color\) to the pit.

`void Remove(int color)`

- Remove a ball of color \(color\) from the pit. If the pit doesn’t contain a ball of the color \(color\), the behavior is undefined.

`bool Constant()`

- Returns true i.f.f. all the balls in the pit are of the same color.

*Can you implement this data structure with less than O(n) memory? (where n is the number of colors)?*

Thanks Ofir Mebel for this really cool riddle!

## Spoiler Alert - Solution Ahead!

All elements in our ball pit data structure are the same i.f.f. the *variance* \(V(x)\) is 0.
Now:

So the ball pit contains elements of the same color i.f.f. \(n\sum{x_i^2} = (\sum{x_i})^2\), and we can implement this data structure with just 3 integer counters: \(n\), \(\sum{x_i}\), \(\sum{x_{i}^{2}}\), as follows:

```
void Insert(int color) {
n++;
sigma_xi += color;
sigma_xi2 += color * color;
}
void Remove(int color) {
n--;
sigma_xi -= color;
sigma_xi2 -= color * color;
}
bool Constant() {
return n * sigma_xi2 == sigma_xi * sigma_xi;
}
```

You can also reason about why this data structure works directly, ignoring the underlying variance calculation. Indeed, obviously if all the balls in the pit are the same color, `Constant`

will return true. To show that `Constant`

cannot return true otherwise, note that the function \(\sum{x_{i}^2}\) strictly gets its minimum (subject to a fixed sum of the \(x_i\)s) when all the \(x_i\) are equal. \(\blacksquare\)