daily_programmer_2015_06_01.py

#http://www.reddit.com/r/dailyprogrammer/comments/3840rp/20150601_challenge_217_easy_lumberjack_pile/
import heapq

class LogPile:
    #Used a min-heap based on value and position to order the locations

    def __init__(self, path):
        self.size = None
        self.logs_to_add = None
        self.logs = {}
        self._read_file(path)

    def _read_file(self, logs_file):

        with open(logs_file, 'r') as f:
            self.size = int(f.readline())
            self.logs_to_add = int(f.readline())

            r = 0
            for line in f:
                c = 0
                for pile in line.split():
                    self.logs[(r,c)] = int(pile)
                    c += 1
                r += 1

    def distribute_logs(self):
        #Handles the Allocation of Logs using a Min-heap should give O(k*log(n)) for k = number of logs to add

        heap = [(v,k) for k,v in self.logs.iteritems()]
        heapq.heapify(heap)

        size, pos = heapq.heappop(heap)

        while self.logs_to_add:
            smallest_pile = (size + 1, pos)
            self.logs_to_add -= 1
            if self.logs_to_add == 0:
                heapq.heappush(heap, smallest_pile)
            else:
                size, pos = heapq.heappushpop(heap, smallest_pile)

        self.logs = {pos : size for size, pos in heap}

    def __repr__(self):
        output = ''
        if self.size == 1:
            return str(self.logs.values()[0])

        else:
            for r in range(self.size - 1):
                output += ' '.join([str(self.logs[(r, c)]) for c in range(self.size)])+'\n'

            return output + ' '.join([str(self.logs[(r + 1, c)]) for c in range(self.size)])

if __name__ == "__main__":
    L = LogPile("C:\\tmp.txt")

    print 'input Logs'
    print L
    print 'expected sum', sum(L.logs.values()) + L.logs_to_add

    L.distribute_logs()

    print
    print 'after distribution'
    print L
    print 'actual sum', sum(L.logs.values()) + L.logs_to_add