Palindromes in a Tree (Codeforces 914E)
Tags: Others, divide and conquer

You are given a tree (a connected acyclic undirected graph) of n vertices. Vertices are numbered from 1 to n and each vertex is assigned a character from a to t.A path in the tree is said to be palindromic if at least one permutation of the labels in the path is a palindrome.For each vertex, output the number of palindromic paths passing through it. Note: The path from vertex u to vertex v is considered to be the same as the path from vertex v to vertex u, and this path will be counted only once for each of the vertices it passes through.

Input: The first line contains an integer n (2 ≤ n ≤ 2·105)  — the number of vertices in the tree.The next n - 1 lines each contain two integers u and v (1  ≤  u, v  ≤  n, u ≠ v) denoting an edge connecting vertex u and vertex v. It is guaranteed that the given graph is a tree.The next line contains a string consisting of n lowercase characters from a to t where the i-th (1 ≤ i ≤ n) character is the label of vertex i in the tree.

Output: Print n integers in a single line, the i-th of which is the number of palindromic paths passing through vertex i in the tree.

Note: In the first sample case, the following paths are palindromic:2 - 3 - 42 - 3 - 54 - 3 - 5Additionally, all paths containing only one vertex are palindromic. Listed below are a few paths in the first sample that are not palindromic:1 - 2 - 31 - 2 - 3 - 41 - 2 - 3 - 5

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

7
6 2
4 3
3 7
5 2
7 2
1 4
afefdfs

Sample Outputs:
1 3 4 3 3

1 4 1 1 2 4 2

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 bfs1(s, x):
    q = [s]
    visit[s] = 1
    y = -1
    for i in q:
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if not visit[j] and h[j] == inf:
                q.append(j)
                visit[j] = 1
                parent[j] = i
            elif not visit[j] and h[j] == x:
                y = j
    return (q, y)

def bfs2(u, x):
    q = [u]
    visit[u] = 1
    color[u] = 0
    for i in q:
        ci = color[i]
        for w in range(s0[i], s0[i + 1]):
            j = G[w]
            if not visit[j] and h[j] >= x:
                q.append(j)
                visit[j] = 1
                color[j] = ci ^ s[j]
                parent[j] = i
    return q
n = int(input())
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
h = [inf] * (n + 1)
parent = [-1] * (n + 1)
(x, s1) = ([], [0] * (n + 3))
for i in range(20):
    visit = [0] * (n + 1)
    (cnt, ma) = ([1] * (n + 1), [0] * (n + 1))
    for s in range(1, n + 1):
        if visit[s] or h[s] ^ inf:
            continue
        (p, y) = bfs1(s, i - 1)
        (z, m) = (s, len(p))
        for j in p[::-1][:-1]:
            if 2 * max(ma[j], m - cnt[j]) <= m:
                z = j
                break
            k = parent[j]
            cnt[k] += cnt[j]
            ma[k] = max(ma[k], cnt[j])
        h[z] = i
        if i:
            x.append(y)
            x.append(z)
            s1[y + 2] += 1
for i in range(2, n + 3):
    s1[i] += s1[i - 1]
H = [0] * n
for i in range(0, 2 * (n - 1), 2):
    j = x[i] + 1
    H[s1[j]] = x[i ^ 1]
    s1[j] += 1
(st, now) = ([x[0]], list(s1))
(u, v) = ([0] * (n + 1), [x[0]])
while st:
    i = st[-1]
    if now[i] == s1[i + 1]:
        st.pop()
    else:
        st.append(H[now[i]])
        now[i] += 1
        u[st[-1]] = len(v)
        v.append(st[-1])
pow2 = [1]
for _ in range(19):
    pow2.append(2 * pow2[-1])
sp = set(pow2 + [0])
s = list(input().rstrip())
s = [0] + [pow2[i - 97] for i in s]
ans = [1] * n
color = [0] * (n + 1)
cnt = [0] * (2 * pow2[19])
for i in range(max(h[1:]), -1, -1):
    visit = [0] * (n + 1)
    dp = [0] * (n + 1)
    for j in range(1, n + 1):
        if h[j] ^ i:
            continue
        p0 = bfs2(j, i)
        for k in p0[1:]:
            cnt[color[k]] += 1
        x = [u[H[k]] for k in range(s1[j], s1[j + 1])]
        x.append(u[j] + len(p0))
        (c, sj) = (0, s[j])
        for k in range(len(x) - 1):
            (l, r) = (x[k], x[k + 1])
            w = [v[z] for z in range(l, r)]
            for z in w:
                cnt[color[z]] -= 1
            for z in w:
                c1 = color[z] ^ sj
                c0 = cnt[c1]
                for p in pow2:
                    c0 += cnt[c1 ^ p]
                dp[z] += c0
                c += c0
            for z in w:
                cnt[color[z]] += 1
        for k in p0[1:]:
            cnt[color[k]] -= 1
            if color[k] ^ sj in sp:
                dp[k] += 1
        dp[j] -= c >> 1
        for k in p0[::-1][:-1]:
            dp[parent[k]] += dp[k]
    for j in range(1, n + 1):
        ans[j - 1] += dp[j]
sys.stdout.write(' '.join(map(str, ans)))

Solution (Python):
import sys, os

def centroid_decomp(coupl):
    n = len(coupl)
    bfs = [n - 1]
    for node in bfs:
        bfs += coupl[node]
        for nei in coupl[node]:
            coupl[nei].remove(node)
    size = [0] * n
    for node in reversed(bfs):
        size[node] = 1 + sum((size[child] for child in coupl[node]))

    def centroid_reroot(root):
        while coupl[root]:
            child = max(coupl[root], key=size.__getitem__)
            if size[child] <= size[root] // 2:
                break
            size[root] -= size[child]
            size[child] += size[root]
            coupl[root].remove(child)
            coupl[child].append(root)
            root = child
        return root
    bfs = [n - 1]
    for node in bfs:
        centroid = centroid_reroot(node)
        bfs += coupl[centroid]
        yield centroid
inp = sys.stdin.buffer.read().split()
ii = 0
n = int(inp[ii])
ii += 1
coupl = [[] for _ in range(n)]
for _ in range(n - 1):
    u = int(inp[ii]) - 1
    ii += 1
    v = int(inp[ii]) - 1
    ii += 1
    coupl[u].append(v)
    coupl[v].append(u)
A = [1 << c - b'a'[0] for c in inp[ii]]
ii += 1
palistates = [0] + [1 << i for i in range(20)]
ans = [0.0] * n
dp = [0.0] * n
val = [0] * n
counter = [0] * (1 << 20)
for centroid in centroid_decomp(coupl):
    bfss = []
    for root in coupl[centroid]:
        bfs = [root]
        for node in bfs:
            bfs += coupl[node]
        bfss.append(bfs)
        for node in bfs:
            val[node] ^= A[node]
            for child in coupl[node]:
                val[child] = val[node]
    entire_bfs = [centroid]
    for bfs in bfss:
        entire_bfs += bfs
    for node in entire_bfs:
        val[node] ^= A[centroid]
        counter[val[node]] += 1
    for bfs in bfss:
        for node in bfs:
            counter[val[node]] -= 1
        for node in bfs:
            v = val[node] ^ A[centroid]
            for p in palistates:
                dp[node] += counter[v ^ p]
        for node in bfs:
            counter[val[node]] += 1
    for node in reversed(entire_bfs):
        dp[node] += sum((dp[child] for child in coupl[node]))
    dp[centroid] += 1
    for p in palistates:
        dp[centroid] += counter[p]
    dp[centroid] //= 2
    for node in entire_bfs:
        ans[node] += dp[node]
        counter[val[node]] = val[node] = 0
        dp[node] = 0.0
os.write(1, b' '.join((str(int(x)).encode('ascii') for x in ans)))

Solution (Python):
import sys
input = sys.stdin.readline

def ctz(x):
    return (x & -x).bit_length() - 1

def shallowest_decomposition_tree(graph, root=0):
    n = len(graph)
    forbidden = [0] * n
    decomposition_tree = [[] for _ in range(n)]
    stacks = [[] for _ in range(n.bit_length())]

    def extract_chain(labels, u):
        while labels:
            label = labels.bit_length() - 1
            labels ^= 2 ** label
            v = stacks[label].pop()
            decomposition_tree[u].append(v)
            u = v
    dfs = [root]
    while dfs:
        u = dfs.pop()
        if u >= 0:
            forbidden[u] = -1
            dfs.append(~u)
            for v in graph[u]:
                if not forbidden[v]:
                    dfs.append(v)
        else:
            u = ~u
            forbidden_once = forbidden_twice = 0
            for v in graph[u]:
                forbidden_twice |= forbidden_once & forbidden[v] + 1
                forbidden_once |= forbidden[v] + 1
            forbidden[u] = forbidden_once | 2 ** forbidden_twice.bit_length() - 1
            label_u = ctz(forbidden[u] + 1)
            stacks[label_u].append(u)
            for v in graph[u]:
                extract_chain(forbidden[v] + 1 & 2 ** label_u - 1, u)
    max_label = (forbidden[root] + 1).bit_length() - 1
    decomposition_root = stacks[max_label].pop()
    extract_chain(forbidden[root] + 1 & 2 ** max_label - 1, decomposition_root)
    return (decomposition_tree, decomposition_root)
n = int(input())
graph = [[] for _ in range(n)]
for _ in range(n - 1):
    (u, v) = [int(x) - 1 for x in input().split()]
    graph[u].append(v)
    graph[v].append(u)
A = [1 << ord(c) - ord('a') for c in input().strip()]
palindromes = [0] + [1 << i for i in range(20)]
(decomposition_tree, decomposition_root) = shallowest_decomposition_tree(graph)
decomp_BFS = [decomposition_root]
for u in decomp_BFS:
    decomp_BFS += decomposition_tree[u]
P = [-1] * n
ans = [1] * n
bitmask = [0] * n
bitmask_count = [0] * 2 ** 20
dp = [0] * n
for decomp_node in decomp_BFS:
    P[decomp_node] = decomp_node
    for v in graph[decomp_node]:
        graph[v].remove(decomp_node)
    bfss = []
    for child in graph[decomp_node]:
        P[child] = decomp_node
        bfs = [child]
        bfss.append(bfs)
        for u in bfs:
            for v in graph[u]:
                if P[v] == -1:
                    P[v] = u
                    bfs.append(v)
    entire_bfs = [decomp_node]
    for bfs in bfss:
        entire_bfs += bfs
    for u in entire_bfs:
        bitmask[u] = bitmask[P[u]] ^ A[u]
        bitmask_count[bitmask[u]] += 1
    for bfs in bfss:
        for u in bfs:
            bitmask_count[bitmask[u]] -= 1
        for u in reversed(bfs):
            base = bitmask[u] ^ A[decomp_node]
            for p in palindromes:
                dp[u] += bitmask_count[base ^ p]
            dp[P[u]] += dp[u]
        for u in bfs:
            bitmask_count[bitmask[u]] += 1
    bitmask_count[A[decomp_node]] -= 1
    for p in palindromes:
        dp[decomp_node] += bitmask_count[p]
    dp[decomp_node] //= 2
    for u in entire_bfs:
        ans[u] += dp[u]
        dp[u] = 0
        P[u] = -1
        bitmask_count[bitmask[u]] = 0
        bitmask[u] = 0
print(*ans)