unilec118_tree

unilec118_tree.py

# динамическое программирование
# возьмём вершину и решим задачу отдельно при условии, что она добавлена в подмножество, и отдельно, что нет
# и выберем вариант, где сумма получилась максимальной
# так как от любой вершины будут начинаться несколько поддеревьев, то будем решать задачи для каждого поддерева

def task_118(edges, value):

    def count_a_if_b_in(a, b, fl):
        # А - для этой вершины решаем задачу
        # B - вершина, откуда мы пришли, -1 - значит ниоткуда
        # fl - включена ли вершина b в подмножество
               
        # будем решать задачи для деревьев, начинающихся из А
        sum_if_a_true = 0 # при условаии, что A в подмножестве
        sum_if_a_false = 0 # при условаии, что A не в подмножестве
        
        # обходим всех соседей А, но не тот лист, откуда пришли
        for i in graph[a] - {b}:
            # всегда делаем расчёт для соседа, с условием, что A не в подмножестве
            sum_if_a_false += count_a_if_b_in(i, a, False)
            # если мы в корне дерева или пришли из вершины не входящей в подмножество, то
            # делаем для соседа расчёт, при условии, что A в подмножестве
            if b == -1 or not fl:
                sum_if_a_true += count_a_if_b_in(i, a, True)
        
        # если это расчёт корневой вершины или мы пришли из вершины не в подмножестве
        # то возвращаем максимум от двух результатов
        if b == -1 or not fl:
            return max(sum_if_a_false, sum_if_a_true + value[a])
        # иначе возвращаем расчёт для случая, если A не в подмножестве
        return sum_if_a_false
    
    # создаём и наполняем граф
    graph = [set() for i in range(0, len(edges) + 1)]
    for i in range(0, len(edges)):
        graph[edges[i][0] - 1].add(edges[i][1] - 1)
        graph[edges[i][1] - 1].add(edges[i][0] - 1)
    
    # вызываем функцию, которая решит задачу, начиная с вершины 1
    return count_a_if_b_in(0, -1, True)

print()

# пример из задачи
edges = [(1, 2), (1, 3), (2, 4)]        
value = [1, 1, 0, 1]
print(task_118(edges, value))

print()

# 0
edges = [(1, 2), (2, 3), (3, 4)]        
value = [-1, -2, -3, -4]
print(task_118(edges, value))

# 1
edges = [(2, 1)]        
value = [1, 0]
print(task_118(edges, value))

# 2
edges = [(2, 1)]        
value = [1, 2]
print(task_118(edges, value))

# 3
edges = [(2, 1), (1, 3)]        
value = [2, 1, 2]
print(task_118(edges, value))

# 4
edges = [(2, 1), (1, 3)]        
value = [4, 1, 2]
print(task_118(edges, value))

# 5
edges = [(2, 1), (1, 3), (3, 4)]        
value = [4, 1, 2, 1]
print(task_118(edges, value))

# 6
edges = [(1, 2), (2, 3), (1, 4), (2, 5), (3, 6)]        
value = [1, 1, 1, 2, 2, 2]
print(task_118(edges, value))

# 7
edges = [(1, 2), (2, 3), (3, 4)]        
value = [-1, -2, 7, -4]
print(task_118(edges, value))

print()