'''ConnectedComponents.py

Compute the 4-connected components of a given image.
Assume the image is given as a list of lists of pixel values.
The output takes the same form.  However, the output pixels
are integers in the range 1 to N, where N is the number of
connected components.

We use a simple implementation of the UNION/FIND ADT, based
on up-trees with no path compression.  Each entry of the
UPTREE_ARRAY is a pair of the form (ii, jj).  If this
pair is equal to the coordinate pair (i, j) of the entry,
that means this entry is the Root of an up-tree.  Otherwise,
(ii, jj) is the pixel location of the Parent of the current
entry.

'''

INPUT_IMAGE = [\
[0, 1, 0, 1],\
[0, 0, 0, 1],\
[1, 1, 1, 1],\
[0, 1, 0, 0]]

def CC(im):
    result = clone(im)
    INIT_UNION_FIND(result)
    for i in range(len(im)):
        for j in range(len(im[0])):
            if j>0:
                v1 = im[i][j] # getPixel
                v2 = im[i][j-1]
                if v1==v2:
                  c1 = FIND(i,j)
                  c2 = FIND(i,j-1)
                  if c1 != c2:
                      UNION(c1, c2)
            if i>0:
                v1 = im[i][j]
                v2 = im[i-1][j]
                if v1==v2:
                  c1 = FIND(i,j)
                  c2 = FIND(i-1,j)
                  if c1 != c2:
                      UNION(c1, c2)
    return RENDER()

def clone(image):
    # Returns a copy of the list of list: image.
    return [row[:] for row in image]

UPTREE_ARRAY = None

def INIT_UNION_FIND(im):
    global UPTREE_ARRAY
    UPTREE_ARRAY = clone(im)
    for i in range(len(im)):
        for j in range(len(im[0])):
            # Make each entry a root:
            UPTREE_ARRAY[i][j] = (i, j)
    print_image(UPTREE_ARRAY)
    
def print_image(im):
    print "-----"
    for row in im:
        print str(row)
        
def FIND(i, j):
    global UPTREE_ARRAY
    (ii, jj) = UPTREE_ARRAY[i][j]
    if ii==i and jj==j:  # At root?
        return (ii, jj)  # We've found the root.
    else: return FIND(ii, jj) # Recursive search from parent.

def UNION(c1, c2):
    (i, j) = c2
    global UPTREE_ARRAY
    UPTREE_ARRAY[i][j] = c1

def RENDER():
    global UPTREE_ARRAY
    n = 0 # components found so far: 0
    print_image(UPTREE_ARRAY)
    out = clone(UPTREE_ARRAY)
    dic = {} # mapping from roots to cc numbers.
    for i in range(len(out)):
        for j in range(len(out[0])):
            (ii, jj) = FIND(i, j)
            if i==ii and j==jj:  # at a root?
                n += 1        # Yes, new component.
                dic[(i,j)]=n  # Register n with the coords.
    # The dictionary is now complete.
    # Use it to assign each pixel its cc number.
    for i in range(len(out)):
        for j in range(len(out[0])):
            out[i][j] = dic[FIND(i, j)]
    return out

print_image(CC(INPUT_IMAGE))