At first, I thought we can use the permutation problem code to solve this problem. However, after some examination, I found the problem is quite different to the permutation one, and we cannot just put the that code into here, at least, we should make some modification to let it work correctly. Realizing this, I turn to find a new idea. That is we use two pointer, one to record the last position of 0 while the other records the last position of 1. Initially their value are both -1, as no 0 or 1 is found by now. We traverse the array from the end to beginning. if we find a 0, do nothing except it's the first 0 we find, set the ptr0 to i. If we find 1, if i > ptr1, set ptr1 = i. Also we compare i with ptr0, if i < ptr0, that means there are some 0 behind the current 1, we exchange the value of the nums[i] and num[ptr0], and compare ptr0 with ptr1, if ptr0 > ptr1, set ptr1 = ptr0, let ptr0--; the above two action is to maintain the ptr0 and ptr1 point to the position of last 0 and 1 respectively. If we find 2, first we compare i with ptr1, if i < ptr1, that means there are some 1(also there may be some 0 ahead of 1) behind 2, from i to ptr1(exclusive), we set nums[j] = nums[j+1] and set nums[ptr1] = 2. Meanwhile ptr0-- and ptr1--. Then we continue the loop. If we don't find 1 behind 2 but find 0 behind 2, the following operations is almost same as the operations in the case of finding 1.
Code:
public class Solution { public void sortColors(int[] nums) { int len = nums.length; //if(len == 0) return; int ptr0 = -1, ptr1 = -1; for(int i = len-1; i >= 0; i--){ int col = nums[i]; if(col == 2){ if(i < ptr1){ for(int j = i; j < ptr1; j++) nums[j] = nums[j+1]; nums[ptr1] = 2; ptr1--; ptr0--; continue; } if(i < ptr0){ nums[i] = 0; nums[ptr0] = 2; ptr0--; continue; } } if(col == 1){ ptr1 = i > ptr1? i:ptr1; if(i < ptr0){ ptr1 = ptr0 > ptr1? ptr0: ptr1; nums[i] = 0; nums[ptr0] = 1; ptr0--; } } if(col == 0) ptr0 = i > ptr0? i:ptr0; } } }