most_balanced_partition.py

test_cases = [
    ([-1, 0, 0, 1, 1, 2], [1, 2, 2, 1, 1, 1]),
    ([-1, 0, 1, 2], [1, 4, 3, 4]),
    ([-1, 0, 0, 0], [10, 11, 10, 10]),
    ([-1, 0], [20, 100]),
    ([-1, 0, 0, 0, 0, 3, 4, 6, 0, 3], [298, 2187, 5054, 266, 1989, 6499, 5450, 2205, 5893, 8095]),
    ([-1, 0, 1, 2, 1, 0, 5, 2, 0, 0], [8475, 6038, 8072, 7298, 5363, 9732, 3786, 5521, 8295, 6186]),
    ([-1, 0, 1, 2, 3, 4, 5, 6, 7, 8], [8618, 5774, 7046, 459, 2279, 2894, 798, 2328, 1011, 2780])
]


def most_balanced_partition(parent, files_size):
    n = len(parent)
    main_sum = [0 for i in range(n)]
    for i in range(n):
        index = i
        while index != -1:
            main_sum[index] += files_size[i]
            index = parent[index]
    sum_without_subtree = main_sum[0] - main_sum[1]
    sum_subtree = main_sum[1]
    min_sum = abs(sum_without_subtree - sum_subtree)
    for i in range(2, n):
        sum_without_subtree = main_sum[0] - main_sum[i]
        sum_subtree = main_sum[i]
        candidate_sum = abs(sum_without_subtree - sum_subtree)
        if min_sum > candidate_sum:
            min_sum = candidate_sum
    return min_sum


if __name__ == '__main__':
    for test in test_cases:
        print(f"min_sum: {most_balanced_partition(parent=test[0], files_size=test[1])}")