SRM 675 Div2 500

ShortestPathWithMagic.py

# -*- coding: utf-8 -*-
import math,string,itertools,fractions,heapq,collections,re,array,bisect

class ShortestPathWithMagic:
    def getTime(self, adj, k):

        n = len(adj)
        d = [[0 for i in range(n)] for j in range(n)]

        # convert string to number
        for row in range(n):
            for col in range(n):
                d[row][col] = int(adj[row][col])

        minCost = [[float('inf') for j in range(k + 1)] for i in range(n)]

        pq = []
        heapq.heappush(pq, (0, 0, 0))

        while len(pq) != 0:
            cost, ind, cnt = heapq.heappop(pq)

            if minCost[ind][cnt] > cost:
                minCost[ind][cnt] = cost

                for i in range(n):
                    if i == ind:
                        continue
                    heapq.heappush(pq, (d[ind][i] + cost, i, cnt))
                    if cnt < k:
                        heapq.heappush(pq, (d[ind][i] * 0.5 + cost, i, cnt + 1))

        ret = float('inf')
        for i in range(k + 1):
            ret = min(ret, minCost[1][i])

        return ret