三色排序

有一个只由0，1，2三种元素构成的整数数组，请使用交换、原地排序而不是使用计数进行排序。

给定一个只含0，1，2的整数数组A及它的大小，请返回排序后的数组。保证数组大小小于等于500。

测试样例：[0,1,1,0,2,2],6返回：[0,0,1,1,2,2]

我的提交

# -*- coding:utf-8 -*-class ThreeColor:    def sortThreeColor(self, A, n):        # write code here        zindex = -1        tindex = n        i = 0        while i < tindex:            if A[i] == 0:                zindex += 1                A[zindex], A[i] = A[i], A[zindex]                i += 1            elif A[i] == 1:                i += 1            else:                tindex -= 1                A[tindex], A[i] = A[i], A[tindex]            print("i = " + str(i) + "  ->  " + str(A))        return Aif __name__ == '__main__':    three = ThreeColor()    print(three.sortThreeColor([1,2,0,2],4))

参考答案

import java.util.*;public class ThreeColor {    public int[] sortThreeColor(int[] A, int n) {        if (A == null || A.length < 2) {            return A;        }        int left = -1;        int index = 0;        int right = A.length;        while (index < right) {            if (A[index] == 0) {                swap(A, ++left, index++);            } else if (A[index] == 2) {                swap(A, index, --right);            } else {                index++;            }        }        return A;    }    public void swap(int[] arr, int index1, int index2) {        int tmp = arr[index1];        arr[index1] = arr[index2];        arr[index2] = tmp;    }}