Editing 3133 – Arbori nr

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

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

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

= Restricții și precizări =

* <code>3 ≤ n ≤ 100.000</code>
* <code>1 ≤ u,v ≤ n</code>, 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 <code>1,3,4</code> sunt frunze, deci au câte <code>0</code> noduri în subarborii cu rădăcinile <code>1, 3</code> respectiv <code>5</code>, <code>2</code> are în subarbore nodul <code>1</code> cu eticheta mai mică decât <code>2</code>, <code>5</code> are în subarbore două noduri cu etichetele mai mici decât <code>5</code>, nodurile <code>1</code> şi <code>4</code>.

== 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>