View source for 0647 – Subarbore2

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

= Date de intrare =
Fișierul de intrare <code>subarbore2IN.txt</code> conține pe prima linie numărul de noduri <code>n</code> și nodul <code>k</code>. 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 <code>subarbore2OUT.txt</code> 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 =

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

= Exemplul 1: =
<code>subarbore2IN.txt</code>
 8 4
 4 3 0 3 2 1 4 1
 -3 -2 4 4 0 -3 3 2 
<code>subarbore2OUT.txt</code>
 3

= Exemplul 2: =
<code>subarbore2IN.txt</code>
 101 2
 4 3 0 3 2 1 4 1
 -3 -2 4 4 0 -3 3 2 
<code>subarbore2OUT.txt</code>
 Datele nu corespund restrictiilor impuse

== Rezolvare ==
<syntaxhighlight lang="python3" line="1">
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()

</syntaxhighlight>