0647 – Subarbore2

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri și un nod k. Fiecare nod al arborelui are asociată o valoare numerică întreagă. Determinați suma valorilor asociate nodurilor din subarborele cu rădăcina în k.

Date de intrare

Fișierul de intrare subarbore2IN.txt conține pe prima linie numărul de noduri n și nodul k. Pe a doua linie se află vectorul de tați al arborelui, valorile fiind separate prin spații. Pe linia a treia se află, în ordine, valorile asociate nodurilor din arbore.

Date de ieșire

Fișierul de ieșire subarbore2OUT.txt va conține pe prima linie valoarea cerută.

Restricții și precizări

• 1 ≤ n ≤ 100
• 1 ≤ k ≤ n
• în vectorul de tați rădăcina este marcată cu 0
• valorile asociate nodurilor din arbore sunt numere întregi din intervalul [-1000,1000]

Exemplul 1:

subarbore2IN.txt

8 4
4 3 0 3 2 1 4 1
-3 -2 4 4 0 -3 3 2 

subarbore2OUT.txt

3

Exemplul 2:

subarbore2IN.txt

101 2
4 3 0 3 2 1 4 1
-3 -2 4 4 0 -3 3 2 

subarbore2OUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

def verificare_restrictii(n, k):
    if not (1 <= n <= 100) or not (1 <= k <= n):
        with open('subarbore2OUT.txt', 'w') as fisier_iesire:
            fisier_iesire.write("Datele nu corespund restrictiilor impuse")
        return False
    return True

def dfs(nod, lista_adiacenta, valori):
    suma_subarbore = valori[nod - 1]  # Corectarea indexării pentru a evita IndexError
    for copil in lista_adiacenta[nod]:
        suma_subarbore += dfs(copil, lista_adiacenta, valori)
    return suma_subarbore

def main():
    # Citirea datelor de intrare
    with open('subarbore2IN.txt', 'r') as fisier:
        n, k = map(int, fisier.readline().split())

    # Verificarea restricțiilor
    if not verificare_restrictii(n, k):
        return

    # Citirea celorlalte date de intrare
    with open('subarbore2IN.txt', 'r') as fisier:
        fisier.readline()  # Trecem peste linia cu n și k
        parinti = list(map(int, fisier.readline().split()))
        valori = list(map(int, fisier.readline().split()))

    # Construirea listei de adiacență
    lista_adiacenta = {i: [] for i in range(1, n + 1)}
    for i in range(n):
        if parinti[i] != 0:
            lista_adiacenta[parinti[i]].append(i + 1)

    # Calcularea sumei valorilor asociate nodurilor din subarborele cu rădăcina în k
    rezultat = dfs(k, lista_adiacenta, valori)

    # Scrierea rezultatului în fișierul de ieșire
    with open('subarbore2OUT.txt', 'w') as fisier_iesire:
        fisier_iesire.write(str(rezultat))

if __name__ == "__main__":
    main()