Alternating Tree (Codeforces 960E)
Tags: DP, divide and conquer

Given a tree with $$$n$$$ nodes numbered from $$$1$$$ to $$$n$$$. Each node $$$i$$$ has an associated value $$$V_i$$$.If the simple path from $$$u_1$$$ to $$$u_m$$$ consists of $$$m$$$ nodes namely $$$u_1 \rightarrow u_2 \rightarrow u_3 \rightarrow \dots u_{m-1} \rightarrow u_{m}$$$, then its alternating function $$$A(u_{1},u_{m})$$$ is defined as $$$A(u_{1},u_{m}) = \sum\limits_{i=1}^{m} (-1)^{i+1} \cdot V_{u_{i}}$$$. A path can also have $$$0$$$ edges, i.e. $$$u_{1}=u_{m}$$$.Compute the sum of alternating functions of all unique simple paths. Note that the paths are directed: two paths are considered different if the starting vertices differ or the ending vertices differ. The answer may be large so compute it modulo $$$10^{9}+7$$$. 

Input: The first line contains an integer $$$n$$$ $$$(2 \leq n \leq 2\cdot10^{5} )$$$ — the number of vertices in the tree.The second line contains $$$n$$$ space-separated integers $$$V_1, V_2, \ldots, V_n$$$ ($$$-10^9\leq V_i \leq 10^9$$$) — values of the nodes.The next $$$n-1$$$ lines each contain two space-separated integers $$$u$$$ and $$$v$$$ $$$(1\leq u, v\leq n, u \neq v)$$$ denoting an edge between vertices $$$u$$$ and $$$v$$$. It is guaranteed that the given graph is a tree.

Output: Print the total sum of alternating functions of all unique simple paths modulo $$$10^{9}+7$$$. 

Note: Consider the first example.A simple path from node $$$1$$$ to node $$$2$$$: $$$1 \rightarrow 2$$$ has alternating function equal to $$$A(1,2) = 1 \cdot (-4)+(-1) \cdot 1 = -5$$$.A simple path from node $$$1$$$ to node $$$3$$$: $$$1 \rightarrow 3$$$ has alternating function equal to $$$A(1,3) = 1 \cdot (-4)+(-1) \cdot 5 = -9$$$.A simple path from node $$$2$$$ to node $$$4$$$: $$$2 \rightarrow 1 \rightarrow 4$$$ has alternating function $$$A(2,4) = 1 \cdot (1)+(-1) \cdot (-4)+1 \cdot (-2) = 3$$$.A simple path from node $$$1$$$ to node $$$1$$$ has a single node $$$1$$$, so $$$A(1,1) = 1 \cdot (-4) = -4$$$.Similarly, $$$A(2, 1) = 5$$$, $$$A(3, 1) = 9$$$, $$$A(4, 2) = 3$$$, $$$A(1, 4) = -2$$$, $$$A(4, 1) = 2$$$, $$$A(2, 2) = 1$$$, $$$A(3, 3) = 5$$$, $$$A(4, 4) = -2$$$, $$$A(3, 4) = 7$$$, $$$A(4, 3) = 7$$$, $$$A(2, 3) = 10$$$, $$$A(3, 2) = 10$$$. So the answer is $$$(-5) + (-9) + 3 + (-4) + 5 + 9 + 3 + (-2) + 2 + 1 + 5 + (-2) + 7 + 7 + 10 + 10 = 40$$$.Similarly $$$A(1,4)=-2, A(2,2)=1, A(2,1)=5, A(2,3)=10, A(3,3)=5, A(3,1)=9, A(3,2)=10, A(3,4)=7, A(4,4)=-2, A(4,1)=2, A(4,2)=3 , A(4,3)=7$$$ which sums upto 40. 

Sample Inputs:
4
-4 1 5 -2
1 2
1 3
1 4

8
-2 6 -4 -4 -9 -3 -7 23
8 2
2 3
1 4
6 5
7 6
4 7
5 8

Sample Outputs:
40

4

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [-1] * (n + 1)
    for k in range(n):
        i = q[k]
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)
n = int(input())
mod = pow(10, 9) + 7
v = [0] + list(map(int, input().split()))
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
(q, dist, parent) = bfs(1)
(dp1, dp2) = ([0] * (n + 1), [0] * (n + 1))
for i in range(1, n + 1):
    if not dist[i] % 2:
        dp1[i] = 1
    else:
        dp2[i] = 1
for i in q[::-1][:-1]:
    j = parent[i]
    dp1[j] += dp1[i]
    dp2[j] += dp2[i]
(o, e) = (dp1[1], dp2[1])
(o0, e0) = (o * (o - 1) // 2 % mod, e * (e - 1) // 2 % mod)
ans = 0
for i in range(1, n + 1):
    (c1, c2) = ([], [])
    di = dist[i]
    for u in range(s0[i], s0[i + 1]):
        j = G[u]
        if di < dist[j]:
            c1.append(dp1[j])
            c2.append(dp2[j])
    if i ^ 1:
        c1.append(o - sum(c1))
        c2.append(e - sum(c2))
        if not dist[i] % 2:
            c1[-1] -= 1
        else:
            c2[-1] -= 1
    d = 1 if not dist[i] % 2 else -1
    c = o0
    for j in c1:
        c -= j * (j - 1) // 2
    ans += d * c % mod * v[i] % mod
    c = e0
    for j in c2:
        c -= j * (j - 1) // 2
    ans -= d * c % mod * v[i] % mod
ans = (2 * ans + sum(v)) % mod
print(ans)

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [-1] * (n + 1)
    for k in range(n):
        i = q[k]
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)
n = int(input())
mod = pow(10, 9) + 7
v = [0] + list(map(int, input().split()))
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
(q, dist, parent) = bfs(1)
(dp1, dp2) = ([0] * (n + 1), [0] * (n + 1))
for i in range(1, n + 1):
    if not dist[i] % 2:
        dp1[i] = 1
    else:
        dp2[i] = 1
for i in q[::-1][:-1]:
    j = parent[i]
    dp1[j] += dp1[i]
    dp2[j] += dp2[i]
(o, e) = (dp1[1], dp2[1])
(o0, e0) = (o * (o - 1) // 2 % mod, e * (e - 1) // 2 % mod)
ans = 0
for i in range(1, n + 1):
    (c1, c2) = ([], [])
    di = dist[i]
    for u in range(s0[i], s0[i + 1]):
        j = G[u]
        if di < dist[j]:
            c1.append(dp1[j])
            c2.append(dp2[j])
    if i ^ 1:
        c1.append(o - sum(c1))
        c2.append(e - sum(c2))
        if not dist[i] % 2:
            c1[-1] -= 1
        else:
            c2[-1] -= 1
    d = 1 if not dist[i] % 2 else -1
    c = o0
    for j in c1:
        c -= j * (j - 1) // 2
    ans += d * c % mod * v[i] % mod
    c = e0
    for j in c2:
        c -= j * (j - 1) // 2
    ans -= d * c % mod * v[i] % mod
ans = (2 * ans + sum(v)) % mod
print(ans)

Solution (Python):
if 1:
    import sys
    input = lambda : sys.stdin.readline().rstrip('\r\n')

    def I():
        return input()

    def II():
        return int(input())

    def MII():
        return map(int, input().split())

    def LI():
        return list(input().split())

    def LII():
        return list(map(int, input().split()))

    def LFI():
        return list(map(float, input().split()))

    def GMI():
        return map(lambda x: int(x) - 1, input().split())

    def LGMI():
        return list(map(lambda x: int(x) - 1, input().split()))
    inf = float('inf')
if 1:
    from io import BytesIO, IOBase
    import math
    import random
    import os
    import bisect
    import typing
    from collections import Counter, defaultdict, deque
    from copy import deepcopy
    from functools import cmp_to_key, lru_cache, reduce
    from heapq import heapify, heappop, heappush, heappushpop, nlargest, nsmallest
    from itertools import accumulate, combinations, permutations, count
    from operator import add, iand, ior, itemgetter, mul, xor
    from string import ascii_lowercase, ascii_uppercase, ascii_letters
    from typing import *
    BUFSIZE = 4096

    class FastIO(IOBase):
        newlines = 0

        def __init__(self, file):
            self._fd = file.fileno()
            self.buffer = BytesIO()
            self.writable = 'x' in file.mode or 'r' not in file.mode
            self.write = self.buffer.write if self.writable else None

        def read(self):
            while True:
                b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                if not b:
                    break
                ptr = self.buffer.tell()
                (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
            self.newlines = 0
            return self.buffer.read()

        def readline(self):
            while self.newlines == 0:
                b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                self.newlines = b.count(b'\n') + (not b)
                ptr = self.buffer.tell()
                (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
            self.newlines -= 1
            return self.buffer.readline()

        def flush(self):
            if self.writable:
                os.write(self._fd, self.buffer.getvalue())
                (self.buffer.truncate(0), self.buffer.seek(0))

    class IOWrapper(IOBase):

        def __init__(self, file):
            self.buffer = FastIO(file)
            self.flush = self.buffer.flush
            self.writable = self.buffer.writable
            self.write = lambda s: self.buffer.write(s.encode('ascii'))
            self.read = lambda : self.buffer.read().decode('ascii')
            self.readline = lambda : self.buffer.readline().decode('ascii')
    sys.stdin = IOWrapper(sys.stdin)
    sys.stdout = IOWrapper(sys.stdout)
(dfs, hashing) = (0, 0)
if dfs:
    from types import GeneratorType

    def bootstrap(f, stack=[]):

        def wrappedfunc(*args, **kwargs):
            if stack:
                return f(*args, **kwargs)
            else:
                to = f(*args, **kwargs)
                while True:
                    if type(to) is GeneratorType:
                        stack.append(to)
                        to = next(to)
                    else:
                        stack.pop()
                        if not stack:
                            break
                        to = stack[-1].send(to)
                return to
        return wrappedfunc
if hashing:
    RANDOM = random.getrandbits(20)

    class Wrapper(int):

        def __init__(self, x):
            int.__init__(x)

        def __hash__(self):
            return super(Wrapper, self).__hash__() ^ RANDOM
read_from_file = 0
if read_from_file:
    file = open('input.txt', 'r').readline().strip()[1:-1]
    fin = open(file, 'r')
    input = lambda : fin.readline().strip()
    output_file = open('output.txt', 'w')

    def fprint(*args, **kwargs):
        print(*args, **kwargs, file=output_file)
mod = 10 ** 9 + 7
n = II()
nums = LII()
path = [[] for _ in range(n)]
for _ in range(n - 1):
    (u, v) = GMI()
    path[u].append(v)
    path[v].append(u)
parent = [-1] * n
dfs_order = []
stack = [0]
while stack:
    u = stack.pop()
    dfs_order.append(u)
    for v in path[u]:
        if parent[u] != v:
            parent[v] = u
            stack.append(v)
dp0 = [0] * n
dp1 = nums[:]
size0 = [0] * n
size1 = [1] * n
for i in reversed(dfs_order):
    if i:
        p = parent[i]
        size0[p] += size1[i]
        size1[p] += size0[i]
        dp0[p] += -dp1[i] + size1[i] * nums[p]
        dp1[p] += -dp0[i] + size0[i] * nums[p]
    dp0[i] %= mod
    dp1[i] %= mod
for p in dfs_order:
    dp0[p] %= mod
    dp1[p] %= mod
    for i in path[p]:
        if parent[i] == p:
            dp0[i] += nums[i] * (size1[p] - size0[i]) - (dp1[p] - size0[i] * nums[p] + dp0[i])
            dp1[i] += nums[i] * (size0[p] - size1[i]) - (dp0[p] - size1[i] * nums[p] + dp1[i])
            size0[i] = size1[p]
            size1[i] = size0[p]
print(sum(dp1) % mod)