Equal Tree Sums (Codeforces 1656E)
Tags: Math, trees

You are given an undirected unrooted tree, i.e. a connected undirected graph without cycles.You must assign a nonzero integer weight to each vertex so that the following is satisfied: if any vertex of the tree is removed, then each of the remaining connected components has the same sum of weights in its vertices.

Input: The input consists of multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10^4$$$) — the number of test cases. Description of the test cases follows.The first line of each test case contains an integer $$$n$$$ ($$$3 \leq n \leq 10^5$$$) — the number of vertices of the tree.The next $$$n-1$$$ lines of each case contain each two integers $$$u, v$$$ ($$$1 \leq u,v \leq n$$$) denoting that there is an edge between vertices $$$u$$$ and $$$v$$$. It is guaranteed that the given edges form a tree.The sum of $$$n$$$ for all test cases is at most $$$10^5$$$.

Output: For each test case, you must output one line with $$$n$$$ space separated integers $$$a_1, a_2, \ldots, a_n$$$, where $$$a_i$$$ is the weight assigned to vertex $$$i$$$. The weights must satisfy $$$-10^5 \leq a_i \leq 10^5$$$ and $$$a_i \neq 0$$$.It can be shown that there always exists a solution satisfying these constraints. If there are multiple possible solutions, output any of them.

Note: In the first case, when removing vertex $$$1$$$ all remaining connected components have sum $$$5$$$ and when removing vertex $$$3$$$ all remaining connected components have sum $$$2$$$. When removing other vertices, there is only one remaining connected component so all remaining connected components have the same sum.

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

Sample Outputs:
-3 5 1 2 2
1 1 1

Solution (Python):
from collections import deque
from sys import stdin, gettrace, stderr
if gettrace():

    def inputi():
        return input()
else:

    def input():
        return next(stdin)[:-1]

    def inputi():
        return stdin.buffer.readline()

def readGraph(n, m):
    adj = [set() for _ in range(n)]
    for _ in range(m):
        (u, v) = map(int, input().split())
        adj[u - 1].add(v - 1)
        adj[v - 1].add(u - 1)
    return adj

def readTree(n):
    return readGraph(n, n - 1)

def solve():
    n = int(input())
    adj = readTree(n)
    q = deque((0,))
    res = [0] * n
    res[0] = len(adj[0])
    while q:
        u = q.popleft()
        for v in adj[u]:
            if res[v] == 0:
                res[v] = len(adj[v]) * (1 if res[u] < 0 else -1)
                q.append(v)
    print(' '.join(map(str, res)))

def main():
    t = int(input())
    for _ in range(t):
        solve()
if __name__ == '__main__':
    main()

Solution (Python):
from collections import deque
from sys import stdin, gettrace, stderr
if gettrace():

    def inputi():
        return input()
else:

    def input():
        return next(stdin)[:-1]

    def inputi():
        return stdin.buffer.readline()

def readGraph(n, m):
    adj = [set() for _ in range(n)]
    for _ in range(m):
        (u, v) = map(int, input().split())
        adj[u - 1].add(v - 1)
        adj[v - 1].add(u - 1)
    return adj

def readTree(n):
    return readGraph(n, n - 1)

def solve():
    n = int(input())
    adj = readTree(n)
    q = deque((0,))
    res = [0] * n
    res[0] = len(adj[0])
    while q:
        u = q.popleft()
        for v in adj[u]:
            if res[v] == 0:
                res[v] = len(adj[v]) * (1 if res[u] < 0 else -1)
                q.append(v)
    print(' '.join(map(str, res)))

def main():
    t = int(input())
    for _ in range(t):
        solve()
if __name__ == '__main__':
    main()

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

def rdvec():
    return [int(i) for i in input().split()]

def rd():
    return int(input())
T = int(input())

def solve():
    n = rd()
    g = [[] for i in range(n)]
    U = [0] * n
    for i in range(n - 1):
        (x, y) = rdvec()
        x -= 1
        y -= 1
        g[x].append(y)
        g[y].append(x)
        U[x] += 1
        U[y] += 1
    todo = [(0, 0, 0)]
    while len(todo) > 0:
        (x, fa, s) = todo.pop()
        if s % 2 == 0:
            U[x] = -U[x]
        for y in g[x]:
            if y != fa:
                todo.append((y, x, s + 1))
    print(*U)
for o in range(T):
    solve()