Imagine you have a hat with N balls which could have any of M colors, what data structure is best if we want to remove ball number X from the hat efficiently. Every time a ball is removed we return which color that ball was. We keep on removing balls one by one (you do not control the position that is requested to be removed) until the hat is empty.

E.g. we have 10 balls with 2 colors, basically it the flatten look of the hat would be: [1,1,2,2,3,3,4,4,5,5] size = 10

A simple solution is to use a flattened array like above, as it yields O(1) removal, every time we remove one we swap it for the last entry, decrement the size but reuse the full array, however that takes up a lot of space O(N * M):

colors: [1,1,2,2,3,3,4,4,5,5] size = 10
remove(5) = 3
colors: [1,1,2,2,3,5,4,4,5,2] size = 9
remove(0) = 1
colors: [5,1,2,2,3,3,4,4,1,2] size = 8
remove(2) = 1
colors: [5,1,4,2,3,3,4,2,1,2] size = 7

Another solution is to use bag (multiset) but then remove isn't super fast.

Any suggestions for data structure optimised for speed and space?