Sorting and array of 0 and 1's in single pass

Question from Princeton.edu's sorting and searching :

Partitioning. Write a static method that sorts a Comparable array that is known to have at most two different values.

Hint: Maintain two pointers, one starting at the left end and moving right, the other starting at the right end and moving left. Maintain the invariant that all elements to the left of the left pointer are equal to the smaller of the two values and all elements to the right of the right pointer are equal to the larger of the two values.

Solution in python:

"""
Sort an array of 0's and 1's

Loop invariant
[0 0 0 0 0 0 |----unknown-----| 1 1 1 1 1 1]
             i                j
All elements < i are 0
All element > j are 1
elements between i and j are unknown.
"""

a1 = [ 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1]
a2 = [ 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
a3 = [ 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 0 ]

def sort_01(a):
    i=0
    j=len(a)-1
    while i < j:
        while i < j and (a[i] < a[j] or a[i]==0): i += 1
        while i < j and a[i] <= a[j]: j -= 1
        if i < j: a[i],a[j] = 0,1
    print(a)

sort_01(a1)
sort_01(a2)
sort_01(a3)