0648 – Subarbore Numărare

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri și un nod k. Determinați câte noduri conține subarborele cu rădăcina în k.

Date de intrare

Fișierul de intrare subarborenumarareIN.txt conține pe prima linie numărul de noduri n și nodul k. 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 subarborenumarareOUT.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

Exemplul 1:

subarborenumarareIN.txt

8 4
4 3 0 3 2 1 4 1

subarborenumarareOUT.txt

5

Exemplul 2:

subarborenumarareIN.txt

101 2
4 3 0 3 2 1 4 1

subarborenumarareOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

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