0652 – Afișare Fii

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri și k noduri distincte din arbore. Afișați fiii fiecăruia dintre cele k noduri.

Date de intrare

Fișierul de intrare afisarefiiIN.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 trei 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 afisarefiiOUT.txt va conține k linii. Linia i va conține numărul de fii ai lui x[i] urmat de un spațiu și de lista fiilor lui x[i], separați și ei printr-un spațiu.Î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 ≤ k ≤ n ≤ 100
• în vectorul de tați rădăcina este marcată cu 0
• dacă un nod x[i] nu are fii, linia corespunzătoare din fișierul de ieșire va contine doar valoarea 0

Exemplul 1:

afisarefiiIN.txt

8
4 3 0 3 2 1 2 1
4
3 4 7 1

afisarefiiOUT.txt

2 2 4 
1 1 
0 
2 6 8 

Exemplul 2:

afisarefiiIN.txt

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

afisarefiiOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

def verifica_restrictii(n, k):
    if not (1 <= k <= n <= 100):
        return "Datele nu corespund restrictiilor impuse"
    return None

def afisare_fii(n, tati, k, noduri_selectate):
    restrictii_mesaj = verifica_restrictii(n, k)
    if restrictii_mesaj:
        return restrictii_mesaj

    rezultat = []

    for nod in noduri_selectate:
        if not (1 <= nod <= n):
            return "Datele nu corespund restrictiilor impuse"

        fii = [i for i, parinte in enumerate(tati, 1) if parinte == nod]
        rezultat.append((len(fii), fii))

    return rezultat

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

# Verificarea restricțiilor și obținerea rezultatului
restrictii_mesaj = verifica_restrictii(n, k)
if restrictii_mesaj:
    with open("afisarefiiOUT.txt", "w") as f:
        f.write(restrictii_mesaj)
else:
    rezultat = afisare_fii(n, tati, k, noduri_selectate)
    with open("afisarefiiOUT.txt", "w") as f:
        if isinstance(rezultat, str):
            f.write(rezultat)
        else:
            for nr_fii, fii in rezultat:
                f.write(f"{nr_fii} {' '.join(map(str, fii))}\n")