Bear and Tree Jumps (Codeforces 771C)
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 os
import sys
from io import BytesIO, IOBase
import heapq as h
from bisect import bisect_left, bisect_right
from types import GeneratorType
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        import os
        self.os = os
        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 = self.os.read(self._fd, max(self.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 = self.os.read(self._fd, max(self.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:
            self.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, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
from collections import defaultdict as dd, deque as dq, Counter as dc
import math, string

def getInts():
    return [int(s) for s in input().split()]

def getInt():
    return int(input())

def getStrs():
    return [s for s in input().split()]

def getStr():
    return input()

def listStr():
    return list(input())

def getMat(n):
    return [getInts() for _ in range(n)]
MOD = 10 ** 9 + 7
''

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

def solve():
    (N, K) = getInts()
    graph = dd(set)
    for i in range(N - 1):
        (A, B) = getInts()
        graph[A].add(B)
        graph[B].add(A)
    dp_count = [[0 for j in range(5)] for i in range(N + 1)]
    dp_total = [0 for j in range(N + 1)]
    global ans
    ans = 0

    @bootstrap
    def dfs(node, parent, depth):
        global ans
        dp_count[node][depth % K] = 1
        dp_total[node] = 1
        for neigh in graph[node]:
            if neigh != parent:
                yield dfs(neigh, node, depth + 1)
                for i in range(K):
                    for j in range(K):
                        diff = (i + j - 2 * depth) % K
                        req = -diff % K
                        ans += req * dp_count[node][i] * dp_count[neigh][j]
                for i in range(K):
                    dp_count[node][i] += dp_count[neigh][i]
                dp_total[node] += dp_total[neigh]
        ans += dp_total[node] * (N - dp_total[node])
        yield
    dfs(1, -1, 0)
    return ans // K
print(solve())

Solution (Python):
import sys
input = sys.stdin.buffer.readline
from collections import deque
(n, k) = map(int, input().split())
tot_dist = 0
vis = [0] * (n + 1)
dist = [[[0] * 2 for j in range(k)] for i in range(n + 1)]
adj = [[] for i in range(n + 1)]
for i in range(n - 1):
    (u, v) = map(int, input().split())
    adj[u].append(v)
    adj[v].append(u)
s = deque([1])
while s:
    c = s[-1]
    if not vis[c]:
        vis[c] = 1
        for ne in adj[c]:
            if not vis[ne]:
                s.append(ne)
    else:
        tot = [0] * k
        sum_dist = 0
        pairable = 0
        for ne in adj[c]:
            tot_dist += pairable * sum([dist[ne][i][0] for i in range(k)]) + sum_dist * sum([dist[ne][i][1] for i in range(k)])
            for i in range(k):
                for j in range(k):
                    tot_dist += (i + j + 2 + k - 1) // k * dist[ne][i][1] * tot[j]
            for i in range(k):
                tot[i] += dist[ne][i][1]
                pairable += dist[ne][i][1]
                sum_dist += dist[ne][i][0]
        for ne in adj[c]:
            for i in range(k):
                for j in range(2):
                    dist[c][i][j] += dist[ne][(i + k - 1) % k][j]
            dist[c][0][0] += dist[ne][k - 1][1]
        dist[c][0][1] += 1
        for i in range(k):
            tot_dist += dist[c][i][0] + (i + k - 1) // k * dist[c][i][1]
        s.pop()
print(tot_dist)

Solution (Python):
if True:
    from io import BytesIO, IOBase
    import math
    import random
    import sys
    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)
    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 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')
(dfs, hashing, timeit) = (False, False, False)
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
if True:
    import time

    def timer(func):

        def wrapper(*args, **kwargs):
            start = time.time()
            result = func(*args, **kwargs)
            end = time.time()
            print(f'Time Used: {round((end - start) * 1000)} ms')
            return result
        return wrapper if timeit else func
(n, k) = MII()
path = [[] for _ in range(n)]
for _ in range(n - 1):
    (u, v) = GMI()
    path[u].append(v)
    path[v].append(u)
stack = [0]
parent = [-1] * n
dfs_order = []
while stack:
    u = stack.pop()
    dfs_order.append(u)
    for v in path[u]:
        if v != parent[u]:
            parent[v] = u
            stack.append(v)
dfs_order.reverse()
tot_depth = [0] * n
tot_count = [1] * n
mod_res = [[0] * n for _ in range(k)]
for j in range(n):
    mod_res[0][j] = 1
ans = 0
for u in dfs_order:
    if u:
        ans += (tot_depth[u] + tot_count[u]) * tot_count[parent[u]] + tot_depth[parent[u]] * tot_count[u]
        tot_depth[parent[u]] += tot_depth[u] + tot_count[u]
        tot_count[parent[u]] += tot_count[u]
        for iu in range(k):
            for ip in range(k):
                v = (-1 - iu - ip) % k
                ans += v * mod_res[iu][u] * mod_res[ip][parent[u]]
        for iu in range(k):
            mod_res[(iu + 1) % k][parent[u]] += mod_res[iu][u]
print(ans // k)