The easy answer first. Count how many reds, yellows and greens there are, then overwrite the array with that many of each - counting sort. O(n) time, O(1) space (three counters), but it takes two passes and it works only because the balls are interchangeable. Say it, then beat it.

One pass: the three-way partition (the Dutch national flag problem). Keep three indices and one invariant that splits the array into four regions:

[0, lo) - all red, finished
[lo, mid) - all yellow, finished
[mid, hi] - unknown, not looked at yet
(hi, n-1] - all green, finished

Start with lo = mid = 0 and hi = n-1, so the unknown region is the whole array. Then keep looking at the ball at mid until the unknown region is empty:

red → swap(lo, mid), then lo++ and mid++
yellow → it is already where it belongs, just mid++
green → swap(mid, hi), then hi-- and mid does not move

lo = mid = 0; hi = n - 1
while mid <= hi:
    if a[mid] == RED:    swap(a, lo, mid); lo += 1; mid += 1
    elif a[mid] == YELLOW: mid += 1
    else:                swap(a, mid, hi); hi -= 1

The one trap - and the thing the question is really testing. On a green swap mid must not advance: the ball that just came back from hi arrived from the unknown region and has never been examined. Advancing past it leaves a red or a yellow stranded in the green band. On a red swap advancing is safe, because the ball coming back from lo is a yellow you have already classified (or lo == mid and it is the same ball).

Why it terminates and why it is correct. Every iteration either increments mid or decrements hi, so the unknown region shrinks by exactly one each time and the loop runs at most n times. The four-region invariant holds after every single step, so when the unknown region is empty the array is sorted by construction.

Cost. O(n) time in a single pass, O(1) extra space - three indices, in place, no allocation. Each ball is looked at once and moved at most once.

Worth adding. This is exactly the partition step of quicksort generalised from two regions to three, which is why the same routine is what makes quicksort survive arrays full of duplicate keys. For two colours only, drop mid and walk two pointers in from both ends.

Array 3 balls