Sorting an array consisting of 0, 1 and 2

sort012.py

import random

def sort012(array):
  	n = len(array)
	p0 = 0 # end of the '0' block
	p2 = n - 1  # beginning of the '2' block

	i = 0
	while i <= p2:
		x = array[i]

		if x == 0:
			if i != p0:
				# we are beyond the '0' block, in the '1' block, so swap
				array[p0] = 0
				array[i] = 1
				p0 += 1
			i += 1

		elif x == 1:
			# otherwise x == 1, so leave it as is
			i += 1

		else:
			assert(x == 2)
			
			# move p2 left until non-2 is found
			while p2 >= 0 and array[p2] == 2:
				p2 -= 1

			if i >= p2:
				# we are right before the beginning of the '2' block, so we are finished
				break

			array[i] = array[p2]
			array[p2] = 2
			p2 -= 1
			# do not i += 1 in this case


def test():
	rand = random.Random()
	for testCase in range(1000):
		array = [rand.randint(0, 2) for x in range(20)]
		sort012(array)
		if array != list(sorted(array)):
			print('Failed:')
			print(array)
			break

test()