Use this to learn the idea, then write your own version.
123 4import collections5import heapq6 7 8class Solution(object):9 def cutOffTree(self, forest):10 """11 :type forest: List[List[int]]12 :rtype: int13 """14 def dot(p1, p2):15 return p1[0]*p2[0]+p1[1]*p2[1]16 17 def minStep(p1, p2):18 min_steps = abs(p1[0]-p2[0])+abs(p1[1]-p2[1])19 closer, detour = [p1], []20 lookup = set()21 while True:22 if not closer: 23 if not detour: 24 return -125 26 min_steps += 227 closer, detour = detour, closer28 i, j = closer.pop()29 if (i, j) == p2:30 return min_steps31 if (i, j) not in lookup:32 lookup.add((i, j))33 for I, J in (i+1, j), (i-1, j), (i, j+1), (i, j-1):34 if 0 <= I < m and 0 <= J < n and forest[I][J] and (I, J) not in lookup:35 is_closer = dot((I-i, J-j), (p2[0]-i, p2[1]-j)) > 036 (closer if is_closer else detour).append((I, J))37 return min_steps38 39 m, n = len(forest), len(forest[0])40 min_heap = []41 for i in xrange(m):42 for j in xrange(n):43 if forest[i][j] > 1:44 heapq.heappush(min_heap, (forest[i][j], (i, j)))45 46 start = (0, 0)47 result = 048 while min_heap:49 tree = heapq.heappop(min_heap)50 step = minStep(start, tree[1])51 if step < 0:52 return -153 result += step54 start = tree[1]55 return result56 57 585960class Solution_TLE(object):61 def cutOffTree(self, forest):62 """63 :type forest: List[List[int]]64 :rtype: int65 """66 def minStep(p1, p2):67 min_steps = 068 lookup = {p1}69 q = collections.deque([p1])70 while q:71 size = len(q)72 for _ in xrange(size):73 (i, j) = q.popleft()74 if (i, j) == p2:75 return min_steps76 for i, j in (i+1, j), (i-1, j), (i, j+1), (i, j-1):77 if not (0 <= i < m and 0 <= j < n and forest[i][j] and (i, j) not in lookup):78 continue79 q.append((i, j))80 lookup.add((i, j))81 min_steps += 182 return -183 84 m, n = len(forest), len(forest[0])85 min_heap = []86 for i in xrange(m):87 for j in xrange(n):88 if forest[i][j] > 1:89 heapq.heappush(min_heap, (forest[i][j], (i, j)))90 91 start = (0, 0)92 result = 093 while min_heap:94 tree = heapq.heappop(min_heap)95 step = minStep(start, tree[1])96 if step < 0:97 return -198 result += step99 start = tree[1]100 return result101 102