LeetCode 2065 - 最大化一张图中的路径价值

from functools import lru_cache
import heapq


class Solution:
    def maximalPathQuality(self, values: List[int], edges: List[List[int]], maxTime: int) -> int:
        n = len(values)

        # 构建邻接表
        adj_list = [[] for _ in range(n)]  # type: list[list[tuple[int, int]]
        for u, v, t in edges:
            adj_list[u].append((v, t))
            adj_list[v].append((u, t))

        # 先计算从节点0到其他节点的最短路径
        # 使用优先队列的版本, 代码简单
        shortest_path = [-1] * n
        pq = [(0, 0)]

        while pq:
            min_time, node = heapq.heappop(pq)
            if shortest_path[node] != -1:
                continue

            shortest_path[node] = min_time
            for v, cost in adj_list[node]:
                heapq.heappush(pq, (cost + min_time, v))

        ans = 0

        def dfs(curr_node: int, remain_time: int, visited: dict[int, int], curr_val: int):
            """
            :param curr_node: 当前节点
            :param remain_time: 剩余时间
            :param visited: 访问过的节点, 使用字典记录访问次数, 之所以不用集合, 是因为可能会访问多次, 无法判断是否需要remove
            :param curr_val: 当前路径累计值
            """

            nonlocal ans
            if remain_time < 0:
                return
            # 如果当前节点能在剩余时间内回到节点0, 说明是有效路径, 此时累积的值是有效的
            if shortest_path[curr_node] <= remain_time:
                ans = max(ans, curr_val)
            else:  # 回不去节点0, 提前返回
                return

            # 使用回溯法暴力枚举所有可能的路径
            for next_node, cost in adj_list[curr_node]:
                visited[next_node] = visited.get(next_node, 0) + 1
                next_val = curr_val + (values[next_node] if visited[next_node] == 1 else 0)
                dfs(next_node, remain_time - cost, visited, next_val)
                visited[next_node] -= 1

        dfs(0, maxTime, {0: 1}, values[0])
        return ans