Python TLE solution, A* search


  • 0
    Y

    I think my algorithm makes sense, using A* algorithm for speed up, still TLE though. Impressed by other passing solutions, bravo~

    class Solution:
        def cutOffTree(self, forest):
            """
            :type forest: List[List[int]]
            :rtype: int
            """
            n = len(forest)
            if n == 0:
                return 0
            m = len(forest[0])
            if m == 0:
                return 0
            positions = [(row, col) for row in range(n) for col in range(m) if forest[row][col] > 1]
            positions = sorted(positions, key=lambda position: forest[position[0]][position[1]])
            current = (0, 0)
            count = 0
            for position in positions:
                step = self.bfs(current, position, forest)
                if step == -1:
                    return -1
                count += step
                current = position
            return count
    
        def bfs(self, start, end, forest):
            n, m = len(forest), len(forest[0])
            if start == end:
                return 0
            visited = [[False] * m for _ in range(n)]
            visited[start[0]][start[1]] = True
            neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1)]
            heap = []
            heapq.heappush(heap, (0, 0, start[0], start[1]))
            while heap:
                f, g, cr, cc = heapq.heappop(heap)
                for dr, dc in neighbors:
                    nr, nc = cr + dr, cc + dc
                    if 0 <= nr < n and 0 <= nc < m and not visited[nr][nc] and forest[nr][nc] > 0:
                        if (nr, nc) == end:
                            return g + 1
                        visited[nr][nc] = True
                        ng = g + 1
                        nf = abs(nr - end[0]) + abs(nc - end[1]) + ng
                        heapq.heappush(heap, (nf, ng, nr, nc))
            return -1
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.