0649 – Subarbori

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri. Determinați câte perechi de noduri neterminale distincte p q din arbore au proprietatea că subarborele cu rădăcina în p și cel cu rădăcina în q au același număr de noduri.

Date de intrare

Fișierul de intrare subarboriIN.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.

Date de ieșire

Fișierul de ieșire subarboriOUT.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
• în vectorul de tați rădăcina este marcată cu 0

Exemplul 1:

subarboriIN.txt

8
4 3 0 3 2 1 2 1

subarboriOUT.txt

1

Explicație

Singura pereche de noduri neterminale pentru care subarborii au același număr de noduri este 2 1. Subarborii cu rădăcina în 2 și 1 au câte 3 noduri fiecare.

Exemplul 2:

subarboriIN.txt

101
4 3 0 3 2 1 2 1

subarboriOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

class NodArbore:
    def __init__(self, valoare):
        self.valoare = valoare
        self.copii = []

def construieste_arbore(tati):
    noduri = [NodArbore(i) for i in range(len(tati))]
    
    for i, tata in enumerate(tati):
        if tata != 0:
            noduri[tata - 1].copii.append(noduri[i])
    
    return noduri[0]  # returnează rădăcina arborelui

def numara_subarbori_cu_acelasi_numar_de_noduri(nod):
    if not nod.copii:
        return 1, {1: 1}  # tuplu: (număr de noduri, dicționar cu numărul de subarbori de acea dimensiune)
    
    total_noduri = 1
    dimensiuni_subarbori = {}
    
    for copil in nod.copii:
        dimensiune_copil, dimensiuni_subarbori_copil = numara_subarbori_cu_acelasi_numar_de_noduri(copil)
        total_noduri += dimensiune_copil
        
        for dimensiune, numar in dimensiuni_subarbori_copil.items():
            dimensiuni_subarbori[dimensiune] = dimensiuni_subarbori.get(dimensiune, 0) + numar
        
    dimensiuni_subarbori[total_noduri] = dimensiuni_subarbori.get(total_noduri, 0) + 1
    
    return total_noduri, dimensiuni_subarbori

def verificare_restricții(n):
    if not (1 <= n <= 100):
        with open("subarboriOUT.txt", "w") as fisier_iesire:
            fisier_iesire.write("Datele nu corespund restrictiilor impuse")
        exit()

def main():
    with open("subarboriIN.txt", "r") as fisier_intrare:
        n = int(fisier_intrare.readline().strip())
        verificare_restricții(n)
        tati = list(map(int, fisier_intrare.readline().strip().split()))

    radacina = construieste_arbore(tati)
    _, dimensiuni_subarbori = numara_subarbori_cu_acelasi_numar_de_noduri(radacina)
    
    rezultat = sum(numar * (numar - 1) // 2 for numar in dimensiuni_subarbori.values())
    
    with open("subarboriOUT.txt", "w") as fisier_iesire:
        fisier_iesire.write(str(rezultat))

if __name__ == "__main__":
    main()