Editing 1666 – Arbrush

== Enunț ==
Eșuând în a-și regăsi adevărata identitate, Brush se refugiază în magicul tărâm al arborilor. Arbotra o sună și îi dă următoarea problemă: se dă un arbore cu <code>N</code> noduri, o rădăcină fixată, și <code>M</code> întrebări de forma: câte perechi neordonate de noduri pot forma, luând noduri doar din subarborele nodului <code>X</code> (inclusiv pe <code>X</code>).

= Cerința =
Ajutați-o pe Brush să răspundă cât mai repede la intrebările lui Arbotra.

= Date de intrare =
Fișierul de intrare <code>arbrushIN.txt</code> va conține pe prima linie trei numere naturale, numărul <code>N</code>, <code>M</code> și <code>Root</code>, separate de un spațiu. Următoarele <code>N-1</code> linii vor conține două numere naturale <code>a</code> și <code>b</code> (între <code>1</code> și <code>N</code>), separate de un spațiu, cu semnificația că există muchie între nodurile <code>a</code> și <code>b</code>. Ultima linie a fișierului de intrare (linia <code>N+1</code>) va conține <code>M</code> numere separate de căte un spațiu, acestea vor semnifica cele <code>M</code> întrebari cerute de Arbotra.

= Date de ieșire =
Fișierul de ieșire <code>arbrushOUT.txt</code> va conține <code>M</code> linii, aflându-se răspunsul la fiecare întrebare cerută (în ordinea cerută în fișierul de intrare).Î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>2 ≤ N, M ≤ 27040</code>
* <code>1 ≤ Root ≤ N</code>
* Uituc de fel, Arbotra poate pune aceeași întrebare de mai multe ori.

= Exemplul 1: =
<code>arbrushIN.txt</code>
 8 4 1
 1 2
 2 3
 2 5
 5 6
 6 7
 6 8
 1 4
 6 4 5 2
<code>arbrushOUT.txt</code>
 3
 0
 6
 15

=== Explicație ===
Pentru prima întrebare, răspunsul este <code>3</code>, deoarece în subarborele nodului <code>6</code> există <code>3</code> noduri: <code>6</code>, <code>7</code>, <code>8</code>, cu care pot forma perechile <code>{6, 7}</code>, <code>{7, 8}</code>, <code>{6, 8}</code>

Pentru a doua întrebare răspunsul este <code>0</code>, deoarece subarborele lui <code>4</code> are doar un nod, pe el înșusi și nu se pot forma perechi de câte două noduri.

Pentru a treia întrebare răspunsul este 6, deoarece în subarborele nodului <code>5</code> există <code>4</code> noduri: <code>5</code>, <code>6</code>, <code>7</code>, <code>8</code>, perechile formate sunt <code>{6, 8}</code>, <code>{5, 6}</code>, <code>{7, 8}</code>, <code>{5, 8}</code>, <code>{5, 7}</code>, <code>{6, 7}</code>

Pentru a patra întrebare răspunsul este <code>15</code>.

== Exemplul 2: ==
<code>arbrushIN.txt</code>
 27041 4 1
 1 2
 2 3
 2 5
 5 6
 6 7
 6 8
 1 4
 6 4 5 2
<code>arbrushOUT.txt</code>
 Datele nu corespund restrictiilor impuse

== Rezolvare ==
<syntaxhighlight lang="python3" line="1">
Nmax = 27040 + 1
Mmax = 27040 + 1

def check_restrictions(N, M, Root):
    if not (2 <= N <= 27040) or not (2 <= M <= 27040) or not (1 <= Root <= N):
        with open("arbrushOUT.txt", "w") as file_out:
            file_out.write("Datele nu corespund restrictiilor impuse\n")
        return False
    return True

def DFS(x):
    global Viz, D, Tree
    Viz[x] = True
    D[x] = 1
    for y in Tree[x]:
        if not Viz[y]:
            DFS(y)
            D[x] += D[y]

def main():
    global N, M, Root, D, Tree, Viz
    with open("arbrushIN.txt", "r") as file_in:
        N, M, Root = map(int, file_in.readline().split())
        if not check_restrictions(N, M, Root):
            return
        
        Tree = [[] for _ in range(Nmax)]
        D = [0] * Nmax
        Viz = [False] * Nmax

        for _ in range(N - 1):
            a, b = map(int, file_in.readline().split())
            Tree[a].append(b)
            Tree[b].append(a)

        DFS(Root)

        questions = list(map(int, file_in.readline().split()))

    with open("arbrushOUT.txt", "w") as file_out:
        for x in questions:
            file_out.write(str((D[x] * (D[x] - 1)) // 2) + '\n')

if __name__ == "__main__":
    main()

</syntaxhighlight>