root/hodgestar/PythonCode/PyramidProblem/pyramid.py

import numpy as np
import numpy.linalg as lg

class Pyramid(object):

    def __init__(self, height, initial):
        self.height = height
        self.initial = initial
        self.solution = None
        self.residual = None

    def _loc(self, row, col):
        return (row*(row-1)/2) + col - 1

    def _below(self, row, col):
        return (row+1, col), (row+1, col+1)

    def solve(self):
        """Solve pyramid puzzle.

        Parameters
        ----------
        height : int
            height of pyramid
        initial: (row, col) -> int
            mapping for initial values
        """
        n = self.height
        n_unknowns = (n+1)*(n)/2
        n_tri_constraints = (n)*(n-1)/2
        n_init_contraints = len(self.initial)
        n_constraints = n_tri_constraints + n_init_contraints

        constraints = np.zeros((n_constraints, n_unknowns))
        rhs = np.zeros((n_constraints, 1))
        constraint = 0
        loc = self._loc
        below = self._below

        # populate triangle constraints
        for row in range(1, n):
            for col in range(1, row+1):
                b1, b2 = below(row, col)
                constraints[constraint, loc(row, col)] = -1
                constraints[constraint, loc(*b1)] = 1
                constraints[constraint, loc(*b2)] = 1
                rhs[constraint] = 0
                constraint += 1

        # populate initial constraints
        for (row, col), val in self.initial.items():
            constraints[constraint, loc(row, col)] = 1
            rhs[constraint] = val
            constraint += 1

        # phew!
        #print constraints
        #print rhs
        self.solution, self.residuals = lg.lstsq(constraints, rhs)[:2]
        #print self.solution

    def __str__(self):
        """Return a string representation of the pyramid."""
        s = []
        s.append("Pyramid (height: %d)" % (self.height,))
        s.append("  [solved: %s]" % ("yes" if self.solution is not None else "no"))
        s.append("  -----")

        if self.solution is not None:
            for row in range(1, self.height+1):
                acol = []
                for col in range(1, row+1):
                    acol.append("%5d" % int(np.round(self.solution[self._loc(row, col)])))
                s.append("  " + "".join(acol))

            s.append("  -----")

            good = True
            for (row, col), val in self.initial.items():
                solval = self.solution[self._loc(row, col)]
                if np.abs(solval - val) > 0.00001:
                    s.append("  !! (%d, %d) != %d (%f)" % (row, col, val, solval))
                    good = False

            if good:
                s.append("  No constraints violated.")

            s.append("  Residual: %f" % np.sum(self.residuals))
            s.append("  -----")

        return "\n".join(s)



def print_pyramid(pyramid):
    """Print a pyramid."""
    print pyramid

if __name__ == "__main__":
    import sys
    if len(sys.argv) < 3:
        print "Usage: pyramid.py <height> <inital value> ..."
        print "e.g. pyramid.py 6 1,1=234 4,2=29 5,2=29 5,4=56 6,1=79 6,4=-69 6,5=125"
        print "e.g. pyramid.py 2 1,1=2 2,1=1"
        sys.exit(1)

    height = int(sys.argv[1])
    initial = {}
    for arg in sys.argv[2:]:
        pos, val = arg.split("=")
        row, col = pos.split(",")
        val = int(val)
        row = int(row)
        col = int(col)
        initial[(row, col)] = val

    pyramid = Pyramid(height, initial)
    pyramid.solve()
    print str(pyramid)