Lintcode 143. Sort Colors II

Given an array of n objects with k different colors (numbered from 1 to k), sort them so that objects of the same color are adjacent, with the colors in the order 1, 2, ... k.

Example

Given colors=[3, 2, 2, 1, 4], k=4, your code should sort colors in-place to [1, 2, 2, 3, 4].

Challenge

A rather straight forward solution is a two-pass algorithm using counting sort. That will cost O(k) extra memory. Can you do it without using extra memory?

Notice

1. You are not suppose to use the library's sort function for this problem.
2. k <= n

Code (Java):

public class Solution {
    /**
     * @param colors: A list of integer
     * @param k: An integer
     * @return: nothing
     */
    public void sortColors2(int[] colors, int k) {
        // write your code here
        if (colors == null || colors.length < 2) {
            return;
        }
        
        sortColors2Helper(colors, 0, colors.length - 1, 1, k);
    }
    
    private void sortColors2Helper(int[] colors, int start, int end, int pStart, int pEnd) {
        if (pStart >= pEnd) {
            return;
        }
        
        int pivot = pStart + (pEnd - pStart) / 2;
        
        int i = start;
        int j = end;
        
        while (i <= j) {
            while (i <= j && colors[i] <= pivot) {
                i++;
            }
            
            while (i <= j && colors[j] > pivot) {
                j--;
            }
            
            if (i <= j) {
                swap(colors, i, j);
                i++;
                j--;
            } 
        }
        
        sortColors2Helper(colors, start, j, pStart, pivot);
        sortColors2Helper(colors, i, end, pivot + 1, pEnd);
    }
    
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}