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\)