Playing at PyData: Pegs

pegs.py

#----------------------------------------------------------------------

def find_non_zero(row, small=1e-5):
    n = len(row)
    for i_col in range(n):
        if abs(row[i_col]) > small:
            return row[i_col]


def normalize(row):
    first_non_zero = find_non_zero(row)
    for (i, val) in enumerate(row):
        row[i] = val / (first_non_zero * 1.0)


def peg_matrix(n, c_prime, matrix):
    for i_row in range(n):
        row = []
        for i_col in range(n):
            if i_row == 0:
                if i_col == 0:
                    row.append(1.5)
                elif i_col == n:
                    row.append(2)
                else:
                    row.append(2)
            elif i_row == n:
                if i_col == 0 or i_col == 1:
                    row.append(1)
                else:
                    row.append(0)
            else:
                if i_col == 0 or i_col == n - i_row + 1 :
                    row.append(1)
                elif i_col < n - i_row + 1:
                    row.append(2)
                else:
                    row.append(0)

        row.append(c_prime[i_row])

        normalize(row)

        matrix.append(row)


def row_reduce(matrix, small=1e-5):
    n = len(matrix)
    for i_row in range(n):
        for j_row in range(i_row + 1, n):
            if matrix[j_row][i_row] > small:
                for i_col in range(i_row, n):
                    matrix[j_row][i_col] = matrix[j_row][i_col] - matrix[i_row][i_col]


def radius(centers):
    """
    solve google pegs problem
    """
    c_prime = []
    for val in centers[1:]:
        c_prime.append(val - centers[0])

    n = len(centers) - 1

    matrix = []

    peg_matrix(n, c_prime, matrix)

    print(matrix)

    # triangularize matrix
    row_reduce(matrix)

    print(matrix)

centers = [4, 30, 50]
centers = [4, 30, 50, 3, 1]
result = radius(centers)
print(result)