0638 – Nivele

Enunț

Într-un arbore cu rădăcină, spunem că rădăcina este pe nivelul 1, fiii rădăcinii pe nivelul 2, fiii fiilor rădăcinii pe nivelul 3, etc.

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri și k noduri din arbore. Determinați pentru fiecare dintre cele k noduri nivelul pe care se află.

Date de intrare

Fișierul de intrare niveleIN.txt conține pe prima linie numărul de noduri n. Pe linia următoare se află vectorul de tați al arborelui, valorile fiind separate prin spații. Linia a treia conține numărul k, iar linia a patra k noduri, x[1], x[2], … , x[k].

Date de ieșire

Fișierul de ieșire niveleOUT.txt va conține k linii. Linia i va conține nivelul pe care se află nodul x[i] în arborele dat

Restricții și precizări

• 1 ≤ k ≤ n ≤ 100
• în vectorul de tați rădăcina este marcată cu 0

Exemplul 1:

niveleIN.txt

7
4 1 7 0 7 7 1 
4
1 3 4 7

niveleOUT.txt

2
4
1
3 

Exemplul 2:

niveleIN.txt

101
4 1 7 0 7 7 1 
4
1 3 4 7

niveleOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

def verifica_restrictii(n, k, tati):
    if not(1 <= k <= n <= 100) or tati.count(0) != 1:
        with open("niveleOUT.txt", "w") as f_out:
            f_out.write("Datele nu corespund restrictiilor impuse")
        return False
    return True

def determina_nivelul(n, tati, k, noduri):
    # Verificare pentru restricții
    if not verifica_restrictii(n, k, tati):
        return

    nivele = [0] * n  # nivelele pentru fiecare nod
    radacina = tati.index(0) + 1  # găsim poziția rădăcinii în vectorul de tați

    def dfs(nod, nivel):
        nivele[nod - 1] = nivel  # nodurile sunt numerotate de la 1, așa că indexăm cu nod-1
        for fiu in range(1, n + 1):
            if tati[fiu - 1] == nod:
                dfs(fiu, nivel + 1)

    dfs(radacina, 1)

    rezultate = [nivele[nod - 1] for nod in noduri]

    # Scrierea rezultatelor în fișierul de ieșire
    with open("niveleOUT.txt", "w") as f_out:
        for rezultat in rezultate:
            f_out.write(str(rezultat) + "\n")

# Citirea datelor de intrare
with open("niveleIN.txt", "r") as f:
    n = int(f.readline().strip())
    tati = list(map(int, f.readline().split()))
    k = int(f.readline().strip())
    noduri = list(map(int, f.readline().split()))

# Determinarea nivelului pentru fiecare nod dat
determina_nivelul(n, tati, k, noduri)