1
# Runtime: 1218 ms (Top 16.46%) | Memory: 23.7 MB (Top 8.86%)
2
class Solution:
3
    def reachableNodes(self, edges: List[List[int]], maxMoves: int, n: int) -> int:
4
        m = defaultdict(list)
5
        for a, b, p in edges:
6
            m[a].append((p, b))
7
            m[b].append((p, a))
8

9
        vis = set()
10
        queue = []
11
        heappush(queue, (0, 0))
12
        edgevis = set()
13
        edgeofl = defaultdict(lambda: 0)
14
        ans = 0
15
        while queue:
16
            # print(queue)
17
            cost, cur = heappop(queue)
18
            vis.add(cur)
19
            for p, nxt in m[cur]:
20
                if p < maxMoves - cost:
21
                    if (cur, nxt) not in edgevis and (nxt, cur) not in edgevis:
22
                        ans += p
23
                        # if nxt in vis:
24
                        # ans -= 1
25
                        edgevis.add((cur, nxt))
26
                        edgevis.add((nxt, cur))
27
                        if nxt not in vis:
28
                            heappush(queue, (cost + p + 1, nxt))
29
                else:
30
                    bal = maxMoves - cost
31
                    if (cur, nxt) in edgevis:
32
                        continue
33
                    if bal <= edgeofl[(cur, nxt)]:
34
                        continue
35
                    if bal + edgeofl[(nxt, cur)] < p:
36
                        ans += bal - edgeofl[(cur, nxt)]
37
                        edgeofl[(cur, nxt)] = bal
38
                    else:
39
                        ans += p - edgeofl[(nxt, cur)] - edgeofl[(cur, nxt)]
40
                        edgevis.add((cur, nxt))
41
                        edgevis.add((nxt, cur))
42
        return ans + len(vis)