Understanding Big-O with a specific example - algorithm

Question

I am working on a rather simple question, to make sure that I understand these concepts.
The question is: there exists an array A of n elements, either being RED, WHITE, or BLUE. Rearrange the array such that all WHITE elements come before all BLUE elements, and all BLUE elements come before all RED elements. Construct an algorithm in O(n) time and O(1) space.
From my understanding, the pseudocode for the solution would be:
numW = numB = 0
for i = 0 to n:
if ARRAY[i] == WHITE:
numW++
else if ARRAY[i] == BLUE:
numB++
for i = 0 to n:
if numW > 0:
ARRAY[i] = WHITE
numW--
else if numB > 0:
ARRAY[i] = BLUE
numB--
else:
ARRAY[i] = RED
I believe it is O(n) because it runs through the loop twice and O(2n) is in O(n). I believe the space is O(1) because it is not dependent on the overall number of elements i.e. there will always be a count for each
Is my understanding correct?

Answer 1

If it's linear time, and your algorithm appears to be, then it's O(n) as you suspect. There's a great summary here: Big-O for Eight Year Olds?

Answer 2

Yes, your solution runs in O(n) time in O(1) space.
Below is my solution which also runs in O(n) time and O(1) space, but also works when we have references to objects, as #kenneth suggested in the comments.
import java.util.Arrays;
import java.util.Random;
import static java.lang.System.out;
class Color{
char c;
Color(char c){
this.c = c;
}
}
public class Solution {
private static void rearrangeColors(Color[] collection){
int ptr = 0;
// move all whites to the left
for(int i=0;i<collection.length;++i){
if(collection[i].c == 'W'){
swap(collection,ptr,i);
ptr++;
}
}
// move all blacks to the left after white
for(int i=ptr;i<collection.length;++i){
if(collection[i].c == 'B'){
swap(collection,ptr,i);
ptr++;
}
}
}
private static void swap(Color[] collection,int ptr1,int ptr2){
Color temp = collection[ptr1];
collection[ptr1] = collection[ptr2];
collection[ptr2] = temp;
}
private static void printColors(Color[] collection){
for(int i=0;i<collection.length;++i){
out.print(collection[i].c + ( i != collection.length - 1 ? "," : ""));
}
out.println();
}
public static void main(String[] args) {
// generate a random collection of 'Color' objects
Random r = new Random();
int array_length = r.nextInt(20) + 1;// to add 1 if in case 0 gets generated
Color[] collection = new Color[array_length];
char[] colors_domain = {'B','W','R'};
for(int i=0;i<collection.length;++i){
collection[i] = new Color(colors_domain[r.nextInt(3)]);
}
// print initial state
printColors(collection);
// rearrange them according to the criteria
rearrangeColors(collection);
// print final state
printColors(collection);
}
}

Answer 3

I won't say this is 100% correct, but a quick test case here did work. If anything, it shows the idea of being able to do it in one pass. Is it faster? Probably not. OP's answer I believe is still the best for this case.
#include <stdio.h>
char temp;
#define SWAP(a,b) { temp = a; a = b; b = temp;}
int main()
{
int n = 10;
char arr[] = "RWBRWBRWBR";
printf("%s\n", arr);
int white = 0;
for(int i=0; i<n; i++)
{
if(arr[i] == 'B')
{
SWAP(arr[i], arr[n-1]);
i--; n--;
}
else if(arr[i] == 'R')
{
SWAP(arr[i], arr[white]);
white++;
}
}
printf("%s\n", arr);
}