0550 - Mere

Cerinţa

Țăranul Ion are în livada sa N pomi, fiecare cu v[i] mere. Între pomi există N-1 cărări, astfel încât între oricare doi pomi să existe un singur drum, alcătuit eventual din mai multe cărări. Pentru că nu și-a plătit ratele la bancă, el este nevoit să vândă o parte dintre pomi. El vrea să adune merele din livadă, dar pentru că nu are foarte mult timp, el va aduna merele doar dintr-o parte din pomi.

Ion pornește din pomul lui preferat, pomul 1, și se deplasează spre unul din vecinii lui. Pentru că nu este foarte inteligent, atunci când Ion se află la un pom, el se va deplasa către pomul vecin care are cele mai multe mere, fără să ia în calcul ceilalți meri din livadă. Dacă doi pomi au același număr de mere, atunci Ion se va deplasa spre pomul cu numărul de ordine mai mic.

Ajutați-l pe Ion să afle câte mere va aduna folosind metoda sa!

Date de intrare

Pe prima linie a fișierului de intrare merein.txt se află un număr natural N, reprezentând numărul de pomi din livadă; pe linia a doua se află N valori v[i], reprezentând numărul de mere din pomul i. Pe următoarele N-1 linii se află câte o pereche de numere x y, cu semnificația că există o cărare de la pomul x la pomul y.

Date de ieșire

Fișierul de ieșire mereout.txt va conține pe prima linie numărul S, reprezentând numărul total de mere culese.

Restricţii şi precizări

• 1 ⩽ n ⩽ 1000
• 1 ⩽ v[i] ⩽ 10000

Exemplul 1

merein.txt

6
3 1 2 3 2 1
1 2
1 3
2 4
2 5
2 6

mereout.txt

Datele de intrare corespund restrictiilor impuse
5

Exemplul 2

merein.txt

65000
3 1 2 3 2 1
1 2
1 3
2 4
2 5
2 6

mereout.txt

Datele de intrare nu corespund restrictiilor impuse

Rezolvare

from collections import defaultdict


def dfs(node, parent):
    global ans, maxv, idx, visit, g, apple
    visit[node] = True
    for neighbour in g[node]:
        if neighbour != parent and not visit[neighbour]:
            dfs(neighbour, node)
            if apple[neighbour] > maxv[node]:
                maxv[node] = apple[neighbour]
                idx[node] = neighbour


def solve():
    global ans, maxv, idx, visit, g, apple
    dfs(1, 0)
    node = 1
    while node:
        ans += apple[node]
        node = idx[node]


if __name__ == "__main__":
    g = defaultdict(list)
    ans = 0

    with open('merein.txt', 'r') as fin:
        N = int(fin.readline().strip())
        if N < 1 or N > 1000:
            with open('mereout.txt', 'w') as fout:
                fout.write("Datele de intrare nu corespund restrictiilor impuse")
            exit(0)
        apple = [0] + list(map(int, fin.readline().split()))
        maxv = [0]*(N+1)
        idx = [0]*(N+1)
        visit = [False]*(N+1)

        for _ in range(N-1):
            u, v = map(int, fin.readline().split())
            g[u].append(v)
            g[v].append(u)

    solve()

    with open('mereout.txt', 'w') as fout:
        fout.write(str(ans) if ans else "Datele de intrare corespund restrictiilor impuse")

Explicatie

Ion va culege merele din pomii 1 3, obținând 3+2=5 mere.