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 N noduri, o rădăcină fixată, și M întrebări de forma: câte perechi neordonate de noduri pot forma, luând noduri doar din subarborele nodului X (inclusiv pe X).

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

Date de ieșire

Fișierul de ieșire arbrushOUT.txt va conține M 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

• 2 ≤ N, M ≤ 27040
• 1 ≤ Root ≤ N
• Uituc de fel, Arbotra poate pune aceeași întrebare de mai multe ori.

Exemplul 1:

arbrushIN.txt

8 4 1
1 2
2 3
2 5
5 6
6 7
6 8
1 4
6 4 5 2

arbrushOUT.txt

3
0
6
15

Explicație

Pentru prima întrebare, răspunsul este 3, deoarece în subarborele nodului 6 există 3 noduri: 6, 7, 8, cu care pot forma perechile {6, 7}, {7, 8}, {6, 8}

Pentru a doua întrebare răspunsul este 0, deoarece subarborele lui 4 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 5 există 4 noduri: 5, 6, 7, 8, perechile formate sunt {6, 8}, {5, 6}, {7, 8}, {5, 8}, {5, 7}, {6, 7}

Pentru a patra întrebare răspunsul este 15.

Exemplul 2:

arbrushIN.txt

27041 4 1
1 2
2 3
2 5
5 6
6 7
6 8
1 4
6 4 5 2

arbrushOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

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()