Vizualizare sursă pentru „0648 – Subarbore Numărare”

= 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>. Determinați câte noduri conține subarborele cu rădăcina în <code>k</code>.

= Date de intrare =
Fișierul de intrare <code>subarborenumarareIN.txt</code> conține pe prima linie numărul de noduri <code>n</code> și nodul <code>k</code>. Pe linia următoare se află vectorul de tați al arborelui, valorile fiind separate prin spații.

= Date de ieșire =
Fișierul de ieșire <code>subarborenumarareOUT.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>

= Exemplul 1: =
<code>subarborenumarareIN.txt</code>
 8 4
 4 3 0 3 2 1 4 1
<code>subarborenumarareOUT.txt</code>
 5

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

== Rezolvare ==
<syntaxhighlight lang="python3" line="1">
def numara_noduri_subarbore(n, k, vector_tati):
    relatii = {}

    for i in range(1, n + 1):
        parinte = vector_tati[i - 1]
        if parinte not in relatii:
            relatii[parinte] = []
        relatii[parinte].append(i)

    def numara_noduri_recursiv(nod):
        numar_noduri = 1
        if nod in relatii:
            for copil in relatii[nod]:
                numar_noduri += numara_noduri_recursiv(copil)
        return numar_noduri

    rezultat = numara_noduri_recursiv(k)

    return rezultat

def verificare_restricții(n, k):
    if not (1 <= n <= 100 and 1 <= k <= n):
        with open("subarborenumarareOUT.txt", "w") as f:
            f.write("Datele nu corespund restrictiilor impuse")
        return False
    return True

# Citire date de intrare
with open("subarborenumarareIN.txt", "r") as f:
    n, k = map(int, f.readline().split())
    vector_tati = list(map(int, f.readline().split()))

# Verificare restricții
if verificare_restricții(n, k):
    # Calcul și scriere rezultat în fișierul de ieșire
    rezultat = numara_noduri_subarbore(n, k, vector_tati)
    with open("subarborenumarareOUT.txt", "w") as f:
        f.write(str(rezultat))

</syntaxhighlight>