Bear and Tree Jumps (Codeforces 791D)
Tags: DP, trees

A tree is an undirected connected graph without cycles. The distance between two vertices is the number of edges in a simple path between them.Limak is a little polar bear. He lives in a tree that consists of n vertices, numbered 1 through n.Limak recently learned how to jump. He can jump from a vertex to any vertex within distance at most k.For a pair of vertices (s, t) we define f(s, t) as the minimum number of jumps Limak needs to get from s to t. Your task is to find the sum of f(s, t) over all pairs of vertices (s, t) such that s < t.

Input: The first line of the input contains two integers n and k (2 ≤ n ≤ 200 000, 1 ≤ k ≤ 5) — the number of vertices in the tree and the maximum allowed jump distance respectively.The next n - 1 lines describe edges in the tree. The i-th of those lines contains two integers ai and bi (1 ≤ ai, bi ≤ n) — the indices on vertices connected with i-th edge.It's guaranteed that the given edges form a tree.

Output: Print one integer, denoting the sum of f(s, t) over all pairs of vertices (s, t) such that s < t.

Note: In the first sample, the given tree has 6 vertices and it's displayed on the drawing below. Limak can jump to any vertex within distance at most 2. For example, from the vertex 5 he can jump to any of vertices: 1, 2 and 4 (well, he can also jump to the vertex 5 itself).  There are  pairs of vertices (s, t) such that s < t. For 5 of those pairs Limak would need two jumps: (1, 6), (3, 4), (3, 5), (3, 6), (5, 6). For other 10 pairs one jump is enough. So, the answer is 5·2 + 10·1 = 20.In the third sample, Limak can jump between every two vertices directly. There are 3 pairs of vertices (s < t), so the answer is 3·1 = 3.

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

13 3
1 2
3 2
4 2
5 2
3 6
10 6
6 7
6 13
5 8
5 9
9 11
11 12

3 5
2 1
3 1

Sample Outputs:
20

114

3

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

class Graph:

    def __init__(self, V, edges=False, graph=False, directed=False, weighted=False, inf=float('inf')):
        self.V = V
        self.directed = directed
        self.weighted = weighted
        self.inf = inf
        if graph:
            self.graph = graph
            self.edges = []
            for i in range(self.V):
                if self.weighted:
                    for (j, d) in self.graph[i]:
                        if self.directed or (not self.directed and i <= j):
                            self.edges.append((i, j, d))
                else:
                    for j in self.graph[i]:
                        if self.directed or (not self.directed and i <= j):
                            self.edges.append((i, j))
        else:
            self.edges = edges
            self.graph = [[] for i in range(self.V)]
            if weighted:
                for (i, j, d) in self.edges:
                    self.graph[i].append((j, d))
                    if not self.directed:
                        self.graph[j].append((i, d))
            else:
                for (i, j) in self.edges:
                    self.graph[i].append(j)
                    if not self.directed:
                        self.graph[j].append(i)

    def SIV_DFS(self, s, bipartite_graph=False, cycle_detection=False, directed_acyclic=False, euler_tour=False, linked_components=False, lowlink=False, parents=False, postorder=False, preorder=False, subtree_size=False, topological_sort=False, unweighted_dist=False, weighted_dist=False):
        seen = [False] * self.V
        finished = [False] * self.V
        if directed_acyclic or cycle_detection or topological_sort:
            dag = True
        if euler_tour:
            et = []
        if linked_components:
            lc = []
        if lowlink:
            order = [None] * self.V
            ll = [None] * self.V
            idx = 0
        if parents or cycle_detection or lowlink or subtree_size:
            ps = [None] * self.V
        if postorder or topological_sort:
            post = []
        if preorder:
            pre = []
        if subtree_size:
            ss = [1] * self.V
        if unweighted_dist or bipartite_graph:
            uwd = [self.inf] * self.V
            uwd[s] = 0
        if weighted_dist:
            wd = [self.inf] * self.V
            wd[s] = 0
        stack = [(s, 0)] if self.weighted else [s]
        while stack:
            if self.weighted:
                (x, d) = stack.pop()
            else:
                x = stack.pop()
            if not seen[x]:
                seen[x] = True
                stack.append((x, d) if self.weighted else x)
                if euler_tour:
                    et.append(x)
                if linked_components:
                    lc.append(x)
                if lowlink:
                    order[x] = idx
                    ll[x] = idx
                    idx += 1
                if preorder:
                    pre.append(x)
                for y in self.graph[x]:
                    if self.weighted:
                        (y, d) = y
                    if not seen[y]:
                        stack.append((y, d) if self.weighted else y)
                        if parents or cycle_detection or lowlink or subtree_size:
                            ps[y] = x
                        if unweighted_dist or bipartite_graph:
                            uwd[y] = uwd[x] + 1
                        if weighted_dist:
                            wd[y] = wd[x] + d
                    elif not finished[y]:
                        if (directed_acyclic or cycle_detection or topological_sort) and dag:
                            dag = False
                            if cycle_detection:
                                cd = (y, x)
            elif not finished[x]:
                finished[x] = True
                if euler_tour:
                    et.append(~x)
                if lowlink:
                    bl = True
                    for y in self.graph[x]:
                        if self.weighted:
                            (y, d) = y
                        if ps[x] == y and bl:
                            bl = False
                            continue
                        ll[x] = min(ll[x], order[y])
                    if x != s:
                        ll[ps[x]] = min(ll[ps[x]], ll[x])
                if postorder or topological_sort:
                    post.append(x)
                if subtree_size:
                    for y in self.graph[x]:
                        if self.weighted:
                            (y, d) = y
                        if y == ps[x]:
                            continue
                        ss[x] += ss[y]
        if bipartite_graph:
            bg = [[], []]
            for tpl in self.edges:
                (x, y) = tpl[:2] if self.weighted else tpl
                if uwd[x] == self.inf or uwd[y] == self.inf:
                    continue
                if not uwd[x] % 2 ^ uwd[y] % 2:
                    bg = False
                    break
            else:
                for x in range(self.V):
                    if uwd[x] == self.inf:
                        continue
                    bg[uwd[x] % 2].append(x)
        retu = ()
        if bipartite_graph:
            retu += (bg,)
        if cycle_detection:
            if dag:
                cd = []
            else:
                (y, x) = cd
                cd = self.Route_Restoration(y, x, ps)
            retu += (cd,)
        if directed_acyclic:
            retu += (dag,)
        if euler_tour:
            retu += (et,)
        if linked_components:
            retu += (lc,)
        if lowlink:
            retu = (ll,)
        if parents:
            retu += (ps,)
        if postorder:
            retu += (post,)
        if preorder:
            retu += (pre,)
        if subtree_size:
            retu += (ss,)
        if topological_sort:
            if dag:
                tp_sort = post[::-1]
            else:
                tp_sort = []
            retu += (tp_sort,)
        if unweighted_dist:
            retu += (uwd,)
        if weighted_dist:
            retu += (wd,)
        if len(retu) == 1:
            retu = retu[0]
        return retu
(N, K) = map(int, readline().split())
edges = []
for _ in range(N - 1):
    (x, y) = map(int, readline().split())
    x -= 1
    y -= 1
    edges.append((x, y))
ans = 0
G = Graph(N, edges=edges)
(parents, postorder) = G.SIV_DFS(0, parents=True, postorder=True)
dp_lst = [None] * N
dp_cnt = [None] * N
ans = 0
for x in postorder:
    dp_lst[x] = [0] * K
    dp_lst[x][0] = 1
    dp_cnt[x] = 0
    lst = [0] * K
    cnt = 0
    for y in G.graph[x]:
        if y == parents[x]:
            continue
        lst = [dp_lst[y][(i - 1) % K] for i in range(K)]
        cnt = dp_cnt[y] + dp_lst[y][K - 1]
        ans -= cnt * sum(lst) * 2
        for i in range(K):
            for j in range(K):
                ans -= lst[i] * lst[j] * ((i + j + K - 1) // K)
        for i in range(K):
            dp_lst[x][i] += dp_lst[y][(i - 1) % K]
        dp_cnt[x] += dp_cnt[y] + dp_lst[y][K - 1]
    ans += dp_cnt[x] * sum(dp_lst[x]) * 2
    for i in range(K):
        for j in range(K):
            ans += dp_lst[x][i] * dp_lst[x][j] * ((i + j + K - 1) // K)
ans //= 2
print(ans)

Solution (Python):
import os, io
if os.access('./test.txt', os.F_OK):
    file = io.open('./test.txt')
else:
    file = io.open(0)
vals = list(map(int, file.readline().split()))
n = vals[0]
k = vals[1]
graph = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    vals = list(map(int, file.readline().split()))
    (u, v) = (vals[0], vals[1])
    graph[u].append(v)
    graph[v].append(u)
count_subtree = [[0] * 5 for _ in range(n + 5)]
total_subtree = [0] * (n + 5)

def subtract(a: int, b: int) -> int:
    return (a - b) % k
stk = [(1, 0, -1, 0)]
ans = 0
while stk:
    (u, nxt, fa, depth) = stk.pop()
    if nxt - 1 >= 0:
        v = graph[u][nxt - 1]
        for i in range(k):
            for j in range(k):
                remainder = subtract(i + j, depth << 1)
                need = subtract(k, remainder)
                ans += need * count_subtree[u][i] * count_subtree[v][j]
        for i in range(k):
            count_subtree[u][i] += count_subtree[v][i]
        total_subtree[u] += total_subtree[v]
    else:
        count_subtree[u][depth % k] = total_subtree[u] = 1
    finish = True
    for ii in range(nxt, len(graph[u])):
        v = graph[u][ii]
        if v != fa:
            stk.append((u, ii + 1, fa, depth))
            stk.append((v, 0, u, depth + 1))
            finish = False
            break
    if finish:
        ans += total_subtree[u] * (n - total_subtree[u])
print(ans // k)
file.close()

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

class Graph:

    def __init__(self, V, edges=False, graph=False, directed=False, weighted=False, inf=float('inf')):
        self.V = V
        self.directed = directed
        self.weighted = weighted
        self.inf = inf
        if graph:
            self.graph = graph
            self.edges = []
            for i in range(self.V):
                if self.weighted:
                    for (j, d) in self.graph[i]:
                        if self.directed or (not self.directed and i <= j):
                            self.edges.append((i, j, d))
                else:
                    for j in self.graph[i]:
                        if self.directed or (not self.directed and i <= j):
                            self.edges.append((i, j))
        else:
            self.edges = edges
            self.graph = [[] for i in range(self.V)]
            if weighted:
                for (i, j, d) in self.edges:
                    self.graph[i].append((j, d))
                    if not self.directed:
                        self.graph[j].append((i, d))
            else:
                for (i, j) in self.edges:
                    self.graph[i].append(j)
                    if not self.directed:
                        self.graph[j].append(i)

    def SIV_DFS(self, s, bipartite_graph=False, cycle_detection=False, directed_acyclic=False, euler_tour=False, linked_components=False, lowlink=False, parents=False, postorder=False, preorder=False, subtree_size=False, topological_sort=False, unweighted_dist=False, weighted_dist=False):
        seen = [False] * self.V
        finished = [False] * self.V
        if directed_acyclic or cycle_detection or topological_sort:
            dag = True
        if euler_tour:
            et = []
        if linked_components:
            lc = []
        if lowlink:
            order = [None] * self.V
            ll = [None] * self.V
            idx = 0
        if parents or cycle_detection or lowlink or subtree_size:
            ps = [None] * self.V
        if postorder or topological_sort:
            post = []
        if preorder:
            pre = []
        if subtree_size:
            ss = [1] * self.V
        if unweighted_dist or bipartite_graph:
            uwd = [self.inf] * self.V
            uwd[s] = 0
        if weighted_dist:
            wd = [self.inf] * self.V
            wd[s] = 0
        stack = [(s, 0)] if self.weighted else [s]
        while stack:
            if self.weighted:
                (x, d) = stack.pop()
            else:
                x = stack.pop()
            if not seen[x]:
                seen[x] = True
                stack.append((x, d) if self.weighted else x)
                if euler_tour:
                    et.append(x)
                if linked_components:
                    lc.append(x)
                if lowlink:
                    order[x] = idx
                    ll[x] = idx
                    idx += 1
                if preorder:
                    pre.append(x)
                for y in self.graph[x]:
                    if self.weighted:
                        (y, d) = y
                    if not seen[y]:
                        stack.append((y, d) if self.weighted else y)
                        if parents or cycle_detection or lowlink or subtree_size:
                            ps[y] = x
                        if unweighted_dist or bipartite_graph:
                            uwd[y] = uwd[x] + 1
                        if weighted_dist:
                            wd[y] = wd[x] + d
                    elif not finished[y]:
                        if (directed_acyclic or cycle_detection or topological_sort) and dag:
                            dag = False
                            if cycle_detection:
                                cd = (y, x)
            elif not finished[x]:
                finished[x] = True
                if euler_tour:
                    et.append(~x)
                if lowlink:
                    bl = True
                    for y in self.graph[x]:
                        if self.weighted:
                            (y, d) = y
                        if ps[x] == y and bl:
                            bl = False
                            continue
                        ll[x] = min(ll[x], order[y])
                    if x != s:
                        ll[ps[x]] = min(ll[ps[x]], ll[x])
                if postorder or topological_sort:
                    post.append(x)
                if subtree_size:
                    for y in self.graph[x]:
                        if self.weighted:
                            (y, d) = y
                        if y == ps[x]:
                            continue
                        ss[x] += ss[y]
        if bipartite_graph:
            bg = [[], []]
            for tpl in self.edges:
                (x, y) = tpl[:2] if self.weighted else tpl
                if uwd[x] == self.inf or uwd[y] == self.inf:
                    continue
                if not uwd[x] % 2 ^ uwd[y] % 2:
                    bg = False
                    break
            else:
                for x in range(self.V):
                    if uwd[x] == self.inf:
                        continue
                    bg[uwd[x] % 2].append(x)
        retu = ()
        if bipartite_graph:
            retu += (bg,)
        if cycle_detection:
            if dag:
                cd = []
            else:
                (y, x) = cd
                cd = self.Route_Restoration(y, x, ps)
            retu += (cd,)
        if directed_acyclic:
            retu += (dag,)
        if euler_tour:
            retu += (et,)
        if linked_components:
            retu += (lc,)
        if lowlink:
            retu = (ll,)
        if parents:
            retu += (ps,)
        if postorder:
            retu += (post,)
        if preorder:
            retu += (pre,)
        if subtree_size:
            retu += (ss,)
        if topological_sort:
            if dag:
                tp_sort = post[::-1]
            else:
                tp_sort = []
            retu += (tp_sort,)
        if unweighted_dist:
            retu += (uwd,)
        if weighted_dist:
            retu += (wd,)
        if len(retu) == 1:
            retu = retu[0]
        return retu
(N, K) = map(int, readline().split())
edges = []
for _ in range(N - 1):
    (x, y) = map(int, readline().split())
    x -= 1
    y -= 1
    edges.append((x, y))
ans = 0
G = Graph(N, edges=edges)
(parents, postorder) = G.SIV_DFS(0, parents=True, postorder=True)
dp_lst = [None] * N
dp_cnt = [None] * N
ans = 0
for x in postorder:
    dp_lst[x] = [0] * K
    dp_lst[x][0] = 1
    dp_cnt[x] = 0
    lst = [0] * K
    cnt = 0
    for y in G.graph[x]:
        if y == parents[x]:
            continue
        lst = [dp_lst[y][(i - 1) % K] for i in range(K)]
        cnt = dp_cnt[y] + dp_lst[y][K - 1]
        ans -= cnt * sum(lst) * 2
        for i in range(K):
            for j in range(K):
                ans -= lst[i] * lst[j] * ((i + j + K - 1) // K)
        for i in range(K):
            dp_lst[x][i] += dp_lst[y][(i - 1) % K]
        dp_cnt[x] += dp_cnt[y] + dp_lst[y][K - 1]
    ans += dp_cnt[x] * sum(dp_lst[x]) * 2
    for i in range(K):
        for j in range(K):
            ans += dp_lst[x][i] * dp_lst[x][j] * ((i + j + K - 1) // K)
ans //= 2
print(ans)