0651 – Sum Sub Max

Cerința

Se dă vectorul de tați al unui arbore cu rădăcină cu n noduri. Fiecare nod al arborelui are asociată o valoare numerică întreagă. Determinați nodurile p din arbore pentru care suma valorilor asociate nodurilor din subarborele cu rădăcina în p este maximă.

Date de intrare

Fișierul de intrare sumsubmaxIN.txt conține pe prima linie numărul de noduri n. Pe a doua linie se află vectorul de tați al arborelui, valorile fiind separate prin spații. Pe linia a treia se află, în ordine, valorile asociate nodurilor din arbore, separate și ele prin spații.

Date de ieșire

Fișierul de ieșire sumsubmaxOUT.txt va conține, în ordine crescătoare, nodurile p din arbore pentru care suma valorilor asociate nodurilor din subarborele cu rădăcina în p este maximă, separate printr-un spațiu.

Restricții și precizări

• 1 ≤ n ≤ 100
• în vectorul de tați rădăcina este marcată cu 0
• valorile asociate nodurilor din arbore sunt numere întregi din intervalul [-1000,1000]

Exemplul 1

sumsubmaxIN.txt:

8

4 3 0 3 2 1 2 1

-3 2 -7 4 0 3 3 1

sumsubmaxOUT.txt:

2 4

Explicație:

În subarborii cu rădăcina în 2 și 4 suma valorilor asociate nodurilor este 5 și este maximă.

Exemplul 2

sumsubmaxIN.txt:

101

4 3 0 3 2 1 2 1

-3 2 -7 4 0 3 3 1

Output:

Invalid input. Please check the constraints.

Rezolvare

def is_valid_input(n, t, val):
    if not (1 <= n <= 100):
        return False
    
    if len(t) != n + 1 or len(val) != n + 1:
        return False
    
    for i in range(1, n + 1):
        if not (-1000 <= val[i] <= 1000):
            return False
    
    return True

def dfs(k):
    v[k] = 1
    for i in range(1, n + 1):
        if i < len(t) and t[i] == k and not v[i]:
            dfs(i)

# Read input from file
with open("sumsubmaxIN.txt", "r") as fin:
    n = int(fin.readline().strip())
    t = [0] + list(map(int, fin.readline().split()))  # Adjust for 1-based indexing
    val = [0] + list(map(int, fin.readline().split()))  # Adjust for 1-based indexing

    # Check input constraints
if not is_valid_input(n, t, val):
    print("Invalid input. Please check the constraints.")
    exit()

smax = [0] * (n + 1)

# Perform DFS and calculate sum of values for each subtree
for i in range(1, n + 1):
    v = [0] * (n + 1)
    dfs(i)
    for j in range(1, n + 1):
        if j < len(v) and v[j] == 1:
            smax[i] += val[j]

# Find the maximum sum
ma = max(smax)

# Write output to file
with open("sumsubmaxOUT.txt", "w") as fout:
    for i in range(1, n + 1):
        if smax[i] == ma:
            fout.write(f"{i} ")