Editing 1895 - Festivaluri (section)

== Rezolvare ==
<syntaxhighlight lang="python3">
from collections import defaultdict, deque

class InParser:
    def __init__(self, nume):
        self.fin = open(nume, "r")

    def read_int(self):
        return int(self.fin.readline())

    def read_values(self):
        return map(int, self.fin.readline().split())

class OutParser:
    def __init__(self, nume):
        self.fout = open(nume, "w")

    def write(self, s):
        self.fout.write(s)

class Graph:
    def __init__(self, n):
        self.adj_list = defaultdict(list)
        self.n = n

    def add_edge(self, u, v, cost):
        self.adj_list[u].append((v, cost))

def max_festivals(n, m, p, z, r, streets, festivals):
    graph = Graph(n)

    for start, end, cost in streets:
        graph.add_edge(start, end, cost)

    visited = set()
    festivals_set = set(festivals)
    queue = deque([(z, 0, p)])  # (current_intersection, current_energy_spent, remaining_energy)

    cnt = 0
    while queue:
        intersection, energy_spent, remaining_energy = queue.popleft()

        if intersection in festivals_set:
            cnt += 1

        if cnt == r:
            return cnt

        if energy_spent > remaining_energy:
            continue

        for neighbor, cost in graph.adj_list[intersection]:
            new_energy_spent = energy_spent + cost
            if new_energy_spent <= p and neighbor not in visited:
                queue.append((neighbor, new_energy_spent, p - new_energy_spent))
                visited.add(neighbor)

    return cnt

fin = InParser("festivaluri.in")
fout = OutParser("festivaluri.out")

n, m, p, z, r = fin.read_values()
streets = [tuple(fin.read_values()) for _ in range(m)]
festivals = list(fin.read_values())

cnt = max_festivals(n, m, p, z, r, streets, festivals)
fout.write(str(cnt) + "\n")

fout.fout.close()

</syntaxhighlight>