Calculate possible combinations of Android pattern lock

android_lock.py

##############################################################################

# R - Red, B- Blue, G-Green
# We name vertices as follows:
# R B R
# B G B
# R B R

##############################################################################

def is_red(position):
    row, column = position
    return (row % 2) + (column % 2) == 0

def is_blue(position):
    row, column = position
    return (row % 2) + (column % 2) == 1

def is_green(position):
    row, column = position
    return row == 1 and column == 1

def get_grid():
    return [(i, j) for i in range(3) for j in range(3)]


def next_steps(position, current_path, allow_repeat=False):
    """ Return possible next steps, given the curren position and path
    covered until now.
    """
    grid = get_grid()

    # Remove self
    grid.remove(position)

    # Remove already visited
    if not allow_repeat:
        for p in current_path:
            if p in grid:
                grid.remove(p)

    if is_green(position):
        # Nothing to remove
        return grid

    elif is_blue(position):
        # Remove blue on opposite edge
        # only if center has not been already used!
        if not allow_repeat and (1, 1) in current_path:
            return grid
        row, column = position
        if row % 2 == 0:
            row = (row * -1) + 2
        else:
            column = (column * -1) + 2
        if (row, column) in grid:
            grid.remove((row, column))
        return grid

    elif is_red(position):
        # Remove other reds
        row, column = position
        row_ = row * -1 + 2
        column_ = column * -1 + 2
        corners = [(i, j) for i in (row, row_) for j in (column, column_)]
        for p in corners:
            if p in grid:
                # Remove only if point in between has not been used!
                mid_point = ((row + p[0])/2, (column + p[1])/2)
                if not allow_repeat and mid_point not in current_path:
                    grid.remove(p)
                elif allow_repeat:
                    grid.remove(p)
        return grid

def count_paths(current_path, length, allow_repeat=False):
    """ Return paths of given length, starting from given point.

    """
    current_position = current_path[-1]
    if length == len(current_path[1:]) or length == 0:
        return 1
    else:
        possible_positions = next_steps(current_position, current_path, allow_repeat)
        count_pp = [count_paths(current_path + [p], length, allow_repeat) \
                    for p in possible_positions]
        return sum(count_pp)

# Tests ######################################################################

def test_colors():
    grid = get_grid()
    for i, p in enumerate(grid):
        if i in [0, 2, 6, 8]:
            assert(is_red(p) == True)
            assert(is_green(p) == False)
            assert(is_blue(p) == False)
        elif i == 4:
            assert(is_green(p) == True)
            assert(is_blue(p) == False)
            assert(is_red(p) == False)
        else:
            assert(is_blue(p) == True)
            assert(is_green(p) == False)
            assert(is_red(p) == False)

def test_next_steps():
    assert(next_steps((0, 0), []) == [(0, 1), (1, 0), (1, 1), (1, 2), (2, 1)])
    assert(next_steps((0, 0), [(1, 1)]) == [(0, 1), (1, 0), (1, 2), (2, 1), (2, 2)])
    assert(next_steps((1, 1), []) == [(0, 0), (0, 1), (0, 2),
                                      (1, 0), (1, 2),
                                      (2, 0), (2, 1), (2, 2)])
    assert(next_steps((0, 1), []) == [(0, 0), (0, 2),
                                      (1, 0), (1, 1), (1, 2),
                                      (2, 0), (2, 2)])
    assert(next_steps((0, 0), [], True) == [(0, 1), (1, 0), (1, 1), (1, 2), (2, 1)])
    assert(next_steps((0, 0), [(1, 1)], True) == [(0, 1), (1, 0), (1, 1), (1, 2), (2, 1)])
    assert(next_steps((1, 1), [], True) == [(0, 0), (0, 1), (0, 2),
                                            (1, 0), (1, 2),
                                            (2, 0), (2, 1), (2, 2)])
    assert(next_steps((0, 1), [], True) == [(0, 0), (0, 2),
                                            (1, 0), (1, 1), (1, 2),
                                            (2, 0), (2, 2)])

def test_count_paths():
    grid = get_grid()
    for i, p in enumerate(grid):
        for length in range(1, 3):
            count = count_paths([p], length)
            count_r = count_paths([p], length, True)
            # Reds
            if i in [0, 2, 6, 8]:
                if length == 1:
                    assert(count == 5)
                    assert(count_r == 5)
                elif length == 2:
                    # 1g*(8-1) + 4b*(7-1)
                    assert(count == 31)
                    # 1g*(8) + 4b*(7)
                    assert(count_r == 36)
            # Green
            elif i == 4:
                if length == 1:
                    assert(count == 8)
                    assert(count_r == 8)
                elif length == 2:
                    # 4r*(5) + 4b*(7)
                    assert(count == 48)
                    # 4r*(5) + 4b*(7)
                    assert(count_r == 48)
            # Blue
            else:
                if length == 1:
                    assert(count == 7)
                    assert(count_r == 7)
                elif length == 2:
                    # 2b*(7-1) + 2r*(5) + 2r*(5-1) + 1g*(8-1)
                    assert(count == 37)
                    # 2b*(7) + 4r*(5) + 1g*(8)
                    assert(count_r == 42)
def test():
    test_colors()
    test_next_steps()
    test_count_paths()

##############################################################################

if __name__ == '__main__':
    #test()

    all_counts = []
    for length in range(1, 10):
        counts = [count_paths([p], length) for p in get_grid()]
        count = sum(counts)
        print length, count, counts
        all_counts.append(count)
    print sum(all_counts[2:])