Box stacking with the maximum height

boxstacking.py

"""Box stacking with the maximum height.

Box Stacking. You are given a set of n types of rectangular 3-D boxes, where
the i^th box has height h(i), width w(i) and depth d(i) (all real numbers). You
want to create a stack of boxes which is as tall as possible, but you can only
stack a box on top of another box if the dimensions of the 2-D base of the
lower box are each strictly larger than those of the 2-D base of the higher
box. Of course, you can rotate a box so that any side functions as its base. It
is also allowable to use multiple instances of the same type of box.

http://people.csail.mit.edu/bdean/6.046/dp/
"""

class Box:
    def __init__(self, h, w, d):
        self.h = h
        self.w = w
        self.d = d

    def __eq__(self, other):
        return self.__dict__ == other.__dict__

    def __str__(self):
        return str(self.__dict__)

class BoxStacking:
    def __init__(self, B):
        self.B = self.expand(B)
        self.n = len(self.B)
        self.tab = [0 for i in xrange(self.n)]
        #self.traceback[i] has the index of self.B that's stacked on top of i-th box
        self.traceback = [None for i in xrange(self.n)]

    def M(self, i):
        if self.tab[i] != 0: return self.tab[i]
        m = [k for k in range(self.n) if self.B[k].w < self.B[i].w
                and self.B[k].d < self.B[i].d]
        if not m:
            self.tab[i] = self.B[i].h
        else:
            n = max(m, key=self.M)
            self.tab[i] = self.B[i].h + self.M(n)
            self.traceback[i] = n
        return self.tab[i]

    def solve(self):
        bottom_ind = max(range(self.n), key=self.M)
        stack = [bottom_ind]
        cur = self.traceback[bottom_ind]
        while cur != None:
            stack.append(cur)
            cur = self.traceback[cur]
        return (self.M(bottom_ind),
                    [self.B[i] for i in stack])

    def norm_base(self, box):
        if box.w < box.d:
            tmp = box.w
            box.w = box.d
            box.d = tmp
        return box

    def rotate(self, box):
        B = []
        B.append(self.norm_base(box))
        b2 = Box(h=box.w, w=box.h, d=box.d)
        b3 = Box(h=box.d, w=box.w, d=box.h)
        B.append(self.norm_base(b2))
        B.append(self.norm_base(b3))
        return B

    def expand(self, B):
        B2 = []
        for box in B:
            B2.extend(self.rotate(box))
        return B2

def test_simple():
    B = [Box(1,2,3)]
    assert BoxStacking(B).solve() == (4, [Box(1,3,2), Box(3,2,1)])

def test_cubic():
    B = [Box(1,1,1)]
    assert BoxStacking(B).solve() == (1, [Box(1,1,1)])

def test_3cubes():
    B = [Box(1,1,1), Box(2,2,2), Box(3,3,3)]
    assert BoxStacking(B).solve() == (6, [Box(3,3,3), Box(2,2,2), Box(1,1,1)])