1757 – Sec

În timp ce-și bea sortimentul preferat de vin sec, vrăjitorului Arpsod i-a venit în minte o problemă de informatică ce are un enunț cel puțin la fel de sec și anume:

Dându-se un arbore binar cu N noduri și rădăcina în nodul 1, să se răspundă la Q întrebări de forma: “sunt cei doi fii ai nodului X identici?”

Doi fii sunt identici dacă au același număr de subarbori și aceștia sunt identici (mai exact, pentru orice i=1, 2, ..., N subarborele i al primului este identic cu subarborele i al celui de-al doilea).

Cerința

Cunoscându-se arborele, să se răspundă la cele Q întrebări de forma indicată în enunţ.

Date de intrare

Pe prima linie a fișierului sec.in se află numărul natural N, reprezentând numărul de noduri ale arborelui. Următoarele N-1 linii conțin perechi de forma ( x, y ) cu semnificația că există muchie între nodul x și nodul y. Pe a (N+1)-a linie se va afla numărul natural Q, reprezentând numărul de întrebări. Pe următoarele Q linii se va afla câte un număr natural reprezentând eticheta nodului ai cărui fii vor fi analizați.

Date de ieșire

Fișierul sec.out va avea Q linii. Pe fiecare linie va fi scris cuvântul DA dacă cei doi fii sunt identici respectiv NU în caz contrar .

Restricții și precizări

• 1 ≤ N ≤ 200.000
• 1 ≤ Q ≤ 500.000
• Nodurile arborelui sunt etichetate de la 1 la N.
• Rădăcina va fi mereu în nodul 1.
• Fii NU sunt identici dacă unul dintre ei trebuie oglindit / rotit pentru a arăta ca celălalt.
• În cazul în care nodul cerut nu are fii, se va afișa DA (se consideră că sunt doi fii NULL identici)
• Se garantează că pentru 30% din teste 1 ≤ N ≤ 1000 și 1 ≤ Q ≤ 3000

Exemplu:

sec.in

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

sec.out

NU
DA
NU
DA

Explicație

Nodurile 2, 7, 8, 9 au ambii fii identici. Celelalte noduri nu au aceasta proprietate

Rezolvare

from collections import defaultdict, deque
import hashlib
import sys

sys.setrecursionlimit(300000)

def read_input(file):
    with open(file, 'r') as f:
        data = f.read().split()
    index = 0
    N = int(data[index])
    index += 1
    edges = []
    for _ in range(N - 1):
        x = int(data[index])
        y = int(data[index + 1])
        index += 2
        edges.append((x, y))
    Q = int(data[index])
    index += 1
    queries = []
    for _ in range(Q):
        queries.append(int(data[index]))
        index += 1
    return N, edges, Q, queries

def build_tree(N, edges):
    tree = defaultdict(list)
    for x, y in edges:
        tree[x].append(y)
        tree[y].append(x)
    return tree

def dfs(node, parent, tree, hashes):
    children = []
    for child in tree[node]:
        if child != parent:
            children.append(dfs(child, node, tree, hashes))
    
    children.sort()
    subtree_hash = hashlib.sha256(''.join(children).encode()).hexdigest()
    hashes[node] = subtree_hash
    return subtree_hash

def solve(N, edges, Q, queries):
    tree = build_tree(N, edges)
    hashes = {}
    
    dfs(1, -1, tree, hashes)
    
    results = []
    for node in queries:
        children = [child for child in tree[node] if child != node]
        if len(children) < 2:
            results.append("DA")
        else:
            if hashes[children[0]] == hashes[children[1]]:
                results.append("DA")
            else:
                results.append("NU")
    
    return results

def write_output(file, results):
    with open(file, 'w') as f:
        for result in results:
            f.write(result + "\n")

def main():
    N, edges, Q, queries = read_input('sec.in')
    results = solve(N, edges, Q, queries)
    write_output('sec.out', results)

if __name__ == "__main__":
    main()