3133 – Arbori nr

Cerința

Se dă un arbore cu n noduri şi rădăcina r, nodurile fiind etichetate cu numerele de la 1 la n. Se cere să se afle pentru fiecare nod i, câte noduri din subarborele cu rădăcina i au eticheta mai mică decât i.

Date de intrare

Fișierul de intrare arbori_nrIN.txt conține pe prima linie numerele n şi r, iar pe următoarele n-1 linii câte o pereche de numere u şi v, reprezentând faptul că între nodurile cu etichetele u şi v există o muchie.

Date de ieșire

Fișierul de ieșire arbori_nrOUT.txt va conține pe prima linie n numere naturale, al i-lea număr reprezentând numărul de noduri din subarborele cu eticheta i, ce au etichetele mai mici decât i.

Restricții și precizări

• 3 ≤ n ≤ 100.000
• 1 ≤ u,v ≤ n, distincte

Exemplul 1

arbori_nrIN.txt:

5 2

2 3

5 4

5 1

2 5

arbori_nrOUT.txt:

0 1 0 0 2

Explicație:

În arbore 1,3,4 sunt frunze, deci au câte 0 noduri în subarborii cu rădăcinile 1, 3 respectiv 5, 2 are în subarbore nodul 1 cu eticheta mai mică decât 2, 5 are în subarbore două noduri cu etichetele mai mici decât 5, nodurile 1 şi 4.

Exemplul 2

arbori_nrIN.txt:

999999999 2

2 3

5 4

5 1

2 5

Output:

Error: n should be between 3 and 100000 (inclusive).

Rezolvare

<syntaxhighlight lang="python3" line="1"> def validate_conditions(n):

   if not (3 <= n <= 100000):
       print("Error: n should be between 3 and 100000 (inclusive).")
       return False
   
   return True

def update(aib, pos, val):

   while pos <= len(aib):
       aib[pos] += val
       pos += pos & -pos

def query(aib, pos):

   sum_val = 0
   while pos > 0:
       sum_val += aib[pos]
       pos -= pos & -pos
   return sum_val

def dfs1(node, dad, G, aib, aiblant, rez, lant):

   update(aib, node, 1)
   update(aiblant, node, 1)
   rez[node] += query(aib, node - 1)
   for next_node in G[node]:
       if next_node != dad:
           dfs1(next_node, node, G, aib, aiblant, rez, lant)
   update(aiblant, node, -1)
   lant[node] = query(aiblant, node)

def dfs2(node, dad, G, aib, rez):

   update(aib, node, 1)
   rez[node] += query(aib, node - 1)
   for next_node in G[node][::-1]:
       if next_node != dad:
           dfs2(next_node, node, G, aib, rez)

def main():

   with open("arbori_nrIN.txt", "r") as f:
       n, r = map(int, f.readline().split())
       if not validate_conditions(n):
           return
       u, v = 0, 0  # Initialize u and v
       G = [[] for _ in range(n + 1)]
       rez = [0] * (n + 1)
       lant = [0] * (n + 1)
       aib = [0] * (4 * n + 1)
       aiblant = [0] * (4 * n + 1)
       for _ in range(n - 1):
           x, y = map(int, f.readline().split())
           G[x].append(y)
           G[y].append(x)
           # Update u and v based on the first edge
           if u == 0:
               u, v = x, y

   dfs1(r, 0, G, aib, aiblant, rez, lant)

   for i in range(4 * n):
       aib[i] = 0

   dfs2(r, 0, G, aib, rez)

   with open("arbori_nrOUT.txt", "w") as g:
       for i in range(1, n + 1):
           g.write(f"{i - rez[i] + lant[i] - 1} ")

if __name__ == "__main__":

   main()

</syntaxhighlight>