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

<syntaxhighlight lang="python3" line="1"> 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)

1. 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)

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]

1. Find the maximum sum

ma = max(smax)

1. 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} ")

</syntaxhighlight>