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

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