0641 – Subarbore

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri și un nod k. Afișați, în ordine crescătoare, nodurile din subarborele cu rădăcina în k.

Date de intrare

Fișierul de intrare subarboreIN.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 subarboreOUT.txt va conține pe prima linie nodurile din subarbore, în ordine crescătoare, separate 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 ≤ n ≤ 100
• 1 ≤ k ≤ n
• în vectorul de tați rădăcina este marcată cu 0

Exemplul 1:

subarboreIN.txt

7 7
4 1 7 0 7 7 1 

subarboreOUT.txt

3 5 6 7

Exemplul 2:

subarboreIN.txt

101 2
4 1 7 0 7 7 1 

subarboreOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

def verifica_restrictii(n, k):
    if not (1 <= n <= 100 and 1 <= k <= n):
        with open("subarboreOUT.txt", "w") as outfile:
            outfile.write("Datele nu corespund restrictiilor impuse")
        return False
    return True

def dfs(nod, lista_adiacenta, noduri_subarbore):
    noduri_subarbore.append(nod)
    for copil in lista_adiacenta[nod]:
        dfs(copil, lista_adiacenta, noduri_subarbore)

def main():
    # Citirea datelor de intrare
    with open("subarboreIN.txt", "r") as infile:
        n, k = map(int, infile.readline().split())
        parinti = list(map(int, infile.readline().split()))

    # Verificarea restricțiilor
    if not verifica_restrictii(n, k):
        return

    # Construirea listei de adiacență bazată pe vectorul de tați
    lista_adiacenta = {i: [] for i in range(1, n + 1)}
    for i in range(1, n + 1):
        if parinti[i - 1] != 0:
            lista_adiacenta[parinti[i - 1]].append(i)

    # Identificarea și colectarea nodurilor din subarbore
    noduri_subarbore = []
    dfs(k, lista_adiacenta, noduri_subarbore)

    # Sortarea și scrierea rezultatelor în fișierul de ieșire
    noduri_subarbore.sort()
    with open("subarboreOUT.txt", "w") as outfile:
        outfile.write(" ".join(map(str, noduri_subarbore)))

if __name__ == "__main__":
    main()