Editing 0674 - Count Sub

= Cerința =
Considerăm un arbore binar cu <code>n</code> noduri în care fiecare nod este numerotat de la <code>1</code> la <code>n</code> și conține o valoare număr natural. Se dau <code>k</code> noduri din arbore și se cere determinarea, pentru fiecare nod, a numărului de noduri din subarborele cu rădăcina în acel nod.

= Date de intrare =
Fișierul de intrare <code>countsubIN.txt</code> conține pe prima linie numărul <code>n</code>. Fiecare dintre următoarele <code>n</code> linii contine câte <code>3</code> numere <code>X st dr</code>; linia <code>i + 1</code> din fișier conține informatiile despre nodul numerotat cu <code>i</code>: <code>X</code> reprezintă valoare din nod, <code>st</code> reprezintă numărul de ordine al descendentului stâng sau <code>0</code> dacă nodul <code>i</code> nu are descendent stâng, iar <code>dr</code> reprezintă numărul de ordine al descendentului drept sau <code>0</code> dacă nodul <code>i</code> nu are descendent drept.

Pe următoarea linie se află numărul <code>k</code>, iar pe fiecare dintre următoarele <code>k</code> linii se află câte un număr natural cuprins între <code>1</code> și <code>n</code>, reprezentând nodul curent.

= Date de ieșire =
Fișierul de ieșire <code>countsubOUT.txt</code> va conține <code>k</code> linii; fiecare linie va conține numărul de noduri din subarborele cu rădăcina în nodul corespunzător. În cazul în care restricțiile nu sunt îndeplinite, se va afișa mesajul "Datele nu corespund restrictiilor impuse".

= Restricții și precizări =

* <code>1 ≤ n ≤ 1000</code>
* valorile din nodurile arborelui vor fi mai mici sau egale cu <code>1.000.000</code>
* <code>1 ≤ k ≤ 1000</code>

= Exemplul 1: =
<code>countsubIN.txt</code>
 6
 2 3 5
 6 0 6
 1 0 0
 7 1 2
 4 0 0
 10 0 0
 4
 1
 2
 4
 2
<code>countsubOUT.txt</code>
 3
 2
 6
 2

= Exemplul 2: =
<code>countsubIN.txt</code>
 1001
 2 3 5
 6 0 6
 1 0 0
 7 1 2
 4 0 0
 10 0 0
 4
 1
 2
 4
 2
<code>countsubOUT.txt</code>
 Datele nu corespund restrictiilor impuse

= Explicație =
Exemplul corespunde arborelui de mai jos, în care au fost marcate cu albastru valorile din noduri, iar cu roșu numerele de ordine ale nodurilor.Rezolvare<syntaxhighlight lang="python" line="1">
def check_restrictions(n, k):
    if not (1 <= n <= 1000):
        return False
    if not (1 <= k <= 1000):
        return False
    return True

def count(k, s, d):
    if k == 0:
        return 0
    else:
        return 1 + count(s[k], s, d) + count(d[k], s, d)

def main():
    with open("countsubIN.txt", "r") as fin:
        try:
            n = int(fin.readline().strip())
        except ValueError:
            with open("countsubOUT.txt", "w") as fout:
                fout.write("Datele nu corespund restrictiilor impuse\n")
            return
        
        s = [0] * (n + 1)
        d = [0] * (n + 1)
        
        for i in range(1, n + 1):
            line = fin.readline().strip().split()
            if len(line) != 3:
                with open("countsubOUT.txt", "w") as fout:
                    fout.write("Datele nu corespund restrictiilor impuse\n")
                return
            x, s[i], d[i] = map(int, line)
        
        try:
            k = int(fin.readline().strip())
        except ValueError:
            with open("countsubOUT.txt", "w") as fout:
                fout.write("Datele nu corespund restrictiilor impuse\n")
            return
        
        queries = []
        for _ in range(k):
            try:
                queries.append(int(fin.readline().strip()))
            except ValueError:
                with open("countsubOUT.txt", "w") as fout:
                    fout.write("Datele nu corespund restrictiilor impuse\n")
                return
    
    if not check_restrictions(n, k):
        with open("countsubOUT.txt", "w") as fout:
            fout.write("Datele nu corespund restrictiilor impuse\n")
        return

    with open("countsubOUT.txt", "w") as fout:
        for y in queries:
            fout.write(f"{count(y, s, d)}\n")

if __name__ == "__main__":
    main()

</syntaxhighlight>