Number of Restricted Paths From First to Last Node LeetCode solution in Python

1
class Solution:
2
    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:
3

4
        dct_nd = {}
5
        dist_to_n = {}
6
        queue = deque()  # from n-node to 1-node
7
        visited = set()
8

9
        # 1 step: create dictionary with nodes and nodes' distances to n
10

11
        # create dictionary with format (weight, node_to)
12
        # heap will automatically sort weight and node_to in ascending order
13
        for l, r, w in edges:
14
            dct_nd[l] = dct_nd.get(l, []) + [(w, r)]
15
            dct_nd[r] = dct_nd.get(r, []) + [(w, l)]
16

17
        dist_to_n[n] = 0
18
        queue.append(n)
19
        visited.add(n)
20

21
        hpf = dct_nd[
22
            n
23
        ].copy()  # without '.copy()' hpf will be only pointer and dct_nd[n] could change
24
        heapify(hpf)
25
        while hpf:
26
            el_w, el_nd = heappop(hpf)
27
            if el_nd in visited:
28
                continue
29

30
            dist_to_n[el_nd] = el_w
31
            visited.add(el_nd)
32
            queue.append(el_nd)
33

34
            if el_nd == 1:
35
                break  # you don't need to traverse more if you've reached 1-node
36
            # other distances will be more than distance of 1-node
37
            for i_w, i_nd in dct_nd[el_nd]:
38
                if i_nd not in visited:
39
                    heappush(hpf, (el_w + i_w, i_nd))
40
        # step 1: end
41

42
        # add to dictionary one more field: number of routes to n
43
        dist_to_n = {k: [v, 0] for k, v in dist_to_n.items()}
44
        dist_to_n[n] = [dist_to_n[n][0], 1]  # for n-node number of routes = 1
45

46
        # step 2: Dynamic Programming
47
        visited.clear()
48
        while queue:
49
            # start from n and traverse futher and futher
50
            nd_prv = queue.popleft()
51
            visited.add(nd_prv)
52
            for w_cur, nd_cur in dct_nd[nd_prv]:
53
                if (
54
                    nd_cur not in visited
55
                    and nd_cur in dist_to_n.keys()
56
                    and dist_to_n[nd_cur][0] > dist_to_n[nd_prv][0]
57
                ):
58
                    # to current node add number of routes from previous node
59
                    dist_to_n[nd_cur][1] += dist_to_n[nd_prv][1]
60
                    # !!! careful !!! you need to add modulo operation (as said in the task)
61
                    dist_to_n[nd_cur][1] = int(dist_to_n[nd_cur][1] % (1e9 + 7))
62
        # step 2: End
63

64
        return dist_to_n[1][1]  # return number of routes for 1-node