Editing 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 <code>N</code> noduri și rădăcina în nodul <code>1</code>, să se răspundă la <code>Q</code> întrebări de forma: “sunt cei doi fii ai nodului <code>X</code> identici?”

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

= Cerința =
Cunoscându-se arborele, să se răspundă la cele <code>Q</code> întrebări de forma indicată în enunţ.

= Date de intrare =
Pe prima linie a fișierului <code>sec.in</code> se află numărul natural <code>N</code>, reprezentând numărul de noduri ale arborelui. Următoarele <code>N-1</code> linii conțin perechi de forma <code>( x, y )</code> cu semnificația că există muchie între nodul <code>x</code> și nodul <code>y</code>. Pe a <code>(N+1)</code>-a linie se va afla numărul natural <code>Q</code>, reprezentând numărul de întrebări. Pe următoarele <code>Q</code> 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 <code>sec.out</code> va avea <code>Q</code> linii. Pe fiecare linie va fi scris cuvântul <code>DA</code> dacă cei doi fii sunt identici respectiv <code>NU</code> în caz contrar .

= Restricții și precizări =

* <code>1 ≤ N ≤ 200.000</code>
* <code>1 ≤ Q ≤ 500.000</code>
* Nodurile arborelui sunt etichetate de la <code>1</code> la <code>N</code>.
* Rădăcina va fi mereu în nodul <code>1</code>.
* 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 <code>DA</code> (se consideră că sunt doi fii <code>NULL</code> identici)
* Se garantează că pentru <code>30%</code> din teste <code>1 ≤ N ≤ 1000</code> și <code>1 ≤ Q ≤ 3000</code>

= Exemplu: =
<code>sec.in</code>
 9
 1 2
 1 3
 2 4
 2 5
 3 6
 4 8
 5 9
 6 7
 4
 1
 2
 3
 7
<code>sec.out</code>
 NU
 DA
 NU
 DA

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

== Rezolvare ==
<syntaxhighlight lang="python3">
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()
</syntaxhighlight>