Sergey and Subway (Codeforces 1060E)
Tags: DP, trees

Sergey Semyonovich is a mayor of a county city N and he used to spend his days and nights in thoughts of further improvements of Nkers' lives. Unfortunately for him, anything and everything has been done already, and there are no more possible improvements he can think of during the day (he now prefers to sleep at night). However, his assistants have found a solution and they now draw an imaginary city on a paper sheet and suggest the mayor can propose its improvements.Right now he has a map of some imaginary city with $$$n$$$ subway stations. Some stations are directly connected with tunnels in such a way that the whole map is a tree (assistants were short on time and enthusiasm). It means that there exists exactly one simple path between each pair of station. We call a path simple if it uses each tunnel no more than once.One of Sergey Semyonovich's favorite quality objectives is the sum of all pairwise distances between every pair of stations. The distance between two stations is the minimum possible number of tunnels on a path between them.Sergey Semyonovich decided to add new tunnels to the subway map. In particular, he connected any two stations $$$u$$$ and $$$v$$$ that were not connected with a direct tunnel but share a common neighbor, i.e. there exists such a station $$$w$$$ that the original map has a tunnel between $$$u$$$ and $$$w$$$ and a tunnel between $$$w$$$ and $$$v$$$. You are given a task to compute the sum of pairwise distances between all pairs of stations in the new map.

Input: The first line of the input contains a single integer $$$n$$$ ($$$2 \leq n \leq 200\,000$$$) — the number of subway stations in the imaginary city drawn by mayor's assistants. Each of the following $$$n - 1$$$ lines contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$, $$$u_i \ne v_i$$$), meaning the station with these indices are connected with a direct tunnel.It is guaranteed that these $$$n$$$ stations and $$$n - 1$$$ tunnels form a tree.

Output: Print one integer that is equal to the sum of distances between all pairs of stations after Sergey Semyonovich draws new tunnels between all pairs of stations that share a common neighbor in the original map.

Note: In the first sample, in the new map all pairs of stations share a direct connection, so the sum of distances is $$$6$$$.In the second sample, the new map has a direct tunnel between all pairs of stations except for the pair $$$(1, 4)$$$. For these two stations the distance is $$$2$$$.

Sample Inputs:
4
1 2
1 3
1 4

4
1 2
2 3
3 4

Sample Outputs:
6

7

Solution (Python):
import sys
import threading
sys.setrecursionlimit(2097152)

def main():

    def countchildren(graph, vert, memo, pard=None):
        dumi = 0
        for child in graph[vert]:
            if child != pard:
                if len(graph[child]) == 1:
                    memo[child] = 0
                else:
                    memo[child] = countchildren(graph, child, memo, vert)[0]
                dumi += memo[child] + 1
        return (dumi, memo)
    n = int(input())
    neigh = []
    for i in range(n):
        neigh.append([])
    for i in range(n - 1):
        (a, b) = map(int, input().split())
        neigh[a - 1].append(b - 1)
        neigh[b - 1].append(a - 1)
    same = 1
    layer = [0]
    pars = [None]
    j = 0
    while layer != []:
        j += 1
        newlayer = []
        newpars = []
        for i in range(len(layer)):
            for vert in neigh[layer[i]]:
                if vert != pars[i]:
                    newlayer.append(vert)
                    newpars.append(layer[i])
        layer = newlayer
        pars = newpars
        if j % 2 == 0:
            same += len(layer)
    bipartite = same * (n - same)
    info = countchildren(neigh, 0, [None] * n)[1]
    dist = 0
    for guy in info:
        if guy != None:
            dist += (guy + 1) * (n - guy - 1)
    print((dist + bipartite) // 2)
threading.stack_size(134217728)
main_thread = threading.Thread(target=main)
main_thread.start()
main_thread.join()

Solution (Python):
def main():

    def countchildren(graph, vert, memo, pard=None):
        dumi = 0
        for child in graph[vert]:
            if child != pard:
                if len(graph[child]) == 1:
                    memo[child] = 0
                else:
                    memo[child] = countchildren(graph, child, memo, vert)[0]
                dumi += memo[child] + 1
        return (dumi, memo)
    n = int(input())
    neigh = []
    for i in range(n):
        neigh.append([])
    for i in range(n - 1):
        (a, b) = map(int, input().split())
        neigh[a - 1].append(b - 1)
        neigh[b - 1].append(a - 1)
    same = 1
    layer = [0]
    pars = [None]
    j = 0
    while layer != []:
        j += 1
        newlayer = []
        newpars = []
        for i in range(len(layer)):
            for vert in neigh[layer[i]]:
                if vert != pars[i]:
                    newlayer.append(vert)
                    newpars.append(layer[i])
        layer = newlayer
        pars = newpars
        if j % 2 == 0:
            same += len(layer)
    bipartite = same * (n - same)
    info = countchildren(neigh, 0, [None] * n)[1]
    dist = 0
    for guy in info:
        if guy != None:
            dist += (guy + 1) * (n - guy - 1)
    print((dist + bipartite) // 2)
import sys
import threading
sys.setrecursionlimit(2097152)
threading.stack_size(134217728)
main_thread = threading.Thread(target=main)
main_thread.start()
main_thread.join()

Solution (Python):
def main():

    def countchildren(graph, vert, memo, pard=None):
        dumi = 0
        for child in graph[vert]:
            if child != pard:
                if len(graph[child]) == 1:
                    memo[child] = 0
                else:
                    memo[child] = countchildren(graph, child, memo, vert)[0]
                dumi += memo[child] + 1
        return (dumi, memo)
    n = int(input())
    neigh = []
    for i in range(n):
        neigh.append([])
    for i in range(n - 1):
        (a, b) = map(int, input().split())
        neigh[a - 1].append(b - 1)
        neigh[b - 1].append(a - 1)
    same = 1
    layer = [0]
    pars = [None]
    j = 0
    while layer != []:
        j += 1
        newlayer = []
        newpars = []
        for i in range(len(layer)):
            for vert in neigh[layer[i]]:
                if vert != pars[i]:
                    newlayer.append(vert)
                    newpars.append(layer[i])
        layer = newlayer
        pars = newpars
        if j % 2 == 0:
            same += len(layer)
    bipartite = same * (n - same)
    info = countchildren(neigh, 0, [None] * n)[1]
    dist = 0
    for guy in info:
        if guy != None:
            dist += (guy + 1) * (n - guy - 1)
    print((dist + bipartite) // 2)
import sys
import threading
sys.setrecursionlimit(2097152)
threading.stack_size(134217728)
main_thread = threading.Thread(target=main)
main_thread.start()
main_thread.join()