[leetcode]Sort Colors @ Python

原題地址：https://oj.leetcode.com/problems/sort-colors/

題意：

Given an array with n objects colored red, white or blue, sort them so that objects of the same color are adjacent, with the colors in the order red, white and blue.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note:
You are not suppose to use the library's sort function for this problem.

click to show follow up.

Follow up:
A rather straight forward solution is a two-pass algorithm using counting sort.
First, iterate the array counting number of 0's, 1's, and 2's, then overwrite array with total number of 0's, then 1's and followed by 2's.

Could you come up with an one-pass algorithm using only constant space?

解題思路：這道題不允許使用排序庫函數。那么最直觀的解法是：遍歷兩遍數組，第一遍對0，1，2計數，第二遍對數組進行賦值，這樣是可以ac的。但題目的要求是只使用常數空間，而且只能遍歷一遍。那么思路就比較巧妙了。設置兩個頭尾指針，頭指針p0指向的位置是0該放置的位置，尾指針p2指向的位置是2該放置的位置。i用來遍歷整個數組，碰到0把它和p0指向的數交換，碰到2把它和p2指向的數交換，碰到1繼續向后遍歷。有點類似快速排序的分割數組這一步。

代碼：

class Solution:
    # @param A a list of integers
    # @return nothing, sort in place
    # @should learn another algorithm
    def sortColors(self, A):
        if A == []: return 
        p0 = 0; p2 = len(A) - 1; i = 0
        while i <= p2:
            if A[i] == 2:
                A[i], A[p2] = A[p2], A[i]
                p2 -= 1
            elif A[i] == 0:
                A[i], A[p0] = A[p0], A[i]
                p0 += 1
                i += 1
            else:
                i += 1