public void sortColors(int[] nums) {
    // The left side of red are all red color,
    // The left side of white are red and white color.
    int red = 0;
    int white = 0;

    for(int i = 0; i < nums.length; i++) {
        if (nums[i] == 0) {
            swap(nums, red, i);
            red++;
            white++;

            // for the situation that we have put the first white element to the i idx,
            // so we need to put it back to white area(to be the last white element).
            if (nums[i] == 1) {
                swap(nums, white - 1, i);
            }
        }else if(nums[i] == 1) {
            swap(nums, white, i);
            white++;
        }   
    }        
}

private void swap(int[] nums, int left, int right) {
    int cur = nums[left];
    nums[left] = nums[right];
    nums[right] = cur;
}