potter.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

DISCOUNTS = [0, 1, 0.95, 0.9, 0.8, 0.75]


def remain(books, n, m):
    """

    Given a set of books, remove m from each of the first n kinds of
books. Return the new book set.

    """
    new_books = books[:]
    for i in xrange(n):
        new_books[i] -= m
    return new_books


def discount(n, m):
    """ Calculate how much are m sets of n different kinds of books. """
    return 8 * n * DISCOUNTS[n] * m


def potter(books):
    """

    calculate the best price of the book set.

    input example:
    [3, 2, 2, 8] means
    3 first books
    2 second books
    2 third books
    8 4th books

    """

    ## sort the books
    books = sorted(books, reverse=True)

    ## remove zeros in the tail
    if 0 in books:
        books = books[:books.index(0)]

    ## if no books, the best price is 0 for sure
    if not books:
        return 0

    ## only 1 kind of books
    elif len(books) == 1:
        return 8 * books[0]

    ## 2 kinds of books
    elif len(books) == 2:
        return discount(2, min(books)) + potter(remain(books, 2, min(books)))

    ## 3 kinds of ...
    elif len(books) == 3:
        return discount(3, min(books)) + potter(remain(books, 3, min(books)))

    ## otherwise
    return min([discount(n, 1) + potter(remain(books, n, 1))
                for n in xrange(len(books), 3, -1)])


def price(set):
    """

    Firstly count how many different kind of books are there,
    then pass the result to function potter to calculate the best price.

    example input:
    [0,0,3,3,3] means

    2 first books
    3 4th books

    """
    books = [0] * 5
    for n in set:
        books[n] += 1
    return potter(books)


def tests():
    def assert_equals(a, b):
        print a, b
        assert a == b

    assert_equals(price([]), 0)
    assert_equals(price([0]), 8)
    assert_equals(price([0, 1]), 8 * 2 * 0.95)
    assert_equals(price([0, 1, 2]), 8 * 3 * .9)
    assert_equals(price([0, 1, 2, 3]), 8 * 4 * .8)
    assert_equals(price([0, 1, 2, 3, 4]), 8 * 5 * .75)
    assert_equals(price([2, 2]), 8 * 2)
    assert_equals(price([0, 0, 1]), 8 * 2 * 0.95 + 8 * 1)
    assert_equals(price([1, 0, 0]), 8 * 2 * 0.95 + 8 * 1)
    assert_equals(price([0, 0, 1, 1]), 2 * (8 * 2 * 0.95))
    assert_equals(price([0, 0, 1, 1, 2, 2, 3, 4]), 2 * (4 * 8 * .80))
    assert_equals(price([0, 0, 1, 1, 2, 3, 3, 4]), 2 * (8 * 4 * .8))
    assert_equals(price([0, 0, 1, 1, 2, 2, 3, 3, 4, 4]), 2 * (8 * 5 * .75))
    assert_equals(price([0, 0, 0, 0, 0,
                         1, 1, 1, 1, 1,
                         2, 2, 2, 2,
                         3, 3, 3, 3, 3,
                         4, 4, 4, 4]),
                  3 * (8 * 5 * 0.75) + 2 * (8 * 4 * 0.8))

    print potter([20, 20, 20, 20, 4])

tests()