Counting transversal paths in a graph

Question

Given a tree with n vertices, each vertex has a special value C_v. A straight path of length k >= 1 is defined as a sequence of vertices v_1, v_2, ... , v_k such that each two consecutive elements of the sequence are connected by an edge and all vertices v_i are different. The straight path may not contain any edges. In other words, for k = 1, a sequence containing a single vertex is also a straight path. There is a function S defined. For a given straight path v_1, v_2, ... , v_k we get S(v_1, v_2, ... ,v_k) = Cv_1 - Cv_2 + Cv_3 - Cv_4 +... Calculate the sum of the values of the function S for all straight paths in the tree. Since the result may be very large, give its remainder when divided by 10^9 + 7. Paths are treated as directed. For example: paths 1 -> 2 -> 4 and 4 -> 2 -> 1 are treated as two different paths and for each one separately the value of the function S should be taken into account in the result.

My implementation is as follows:

def S(path):
    total, negative_one_pow = 0, 1
    for node in path:
        total += (values[node - 1] * negative_one_pow)
        negative_one_pow *= -1
    return total


def search(graph):
    global total
    for node in range(1, n + 1):
        queue = [(node, [node])]
        visited = set()
        while queue:
            current_node, path = queue.pop(0)
            if current_node in visited:
                continue
            visited.add(current_node)
            total += S(path)
            for neighbor in graph[current_node]:
                queue.append((neighbor, [*path, neighbor]))


n = int(input())
values = list(map(int, input().split()))
graph = {i: [] for i in range(1, n + 1)}
total = 0

for i in range(n - 1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

search(graph)
print(total % 1000000007)

The execution of the code takes too long for bigger graphs. Can you suggest ways to speed up the code?

Answer 1

It is a tree. Therefore all straight paths start somewhere, travel up some distance (maybe 0), hit a peak, then turn around and go down some distance (maybe 0).

First calculate for each node the sum and count of all odd length paths rising to its children, and the sum and count of all even length paths rising to its children. This can be done through dynamic programming.

Armed with that we can calculate for each node the sum of all paths with that node as a peak. (Calculate the sum of all paths that rise to the peak then fall. For each child subtract out the sum of all paths that rose to that child and then fell back to that child.)