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ă.În cazul în care restricțiile nu sunt îndeplinite, se va afișa mesajul "Nu corespunde restricțiilor impuse".

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()