Hotel Construction

import sys, bisect, collections, heapq, functools
input = sys.stdin.readline
from math import factorial as fact

rs = lambda: input().strip()
ri = lambda: int(input())
rmi = lambda: map(int, input().split())
ra = lambda: [int(x) for x in input().split()]

INF = 10**18
MOD = 10**9 + 7

def nCr(n,r):
    return fact(n)//(fact(n-r) * fact(r))

def solve():
    dist = [[INF] * n for _ in range(n)]
    for u, v in A:
        u -= 1
        v -= 1
        dist[u][v] = 1
        dist[v][u] = 1

    #Floyd-Warshall DP to calculate dist between any node i and j
    for k in range(n):
        for i in range(n):
            for j in range(n):
                dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

    #Store frequency of all possible distances between any pair i, j (i != j)
    freq = collections.Counter()
    for i in range(n):
        for j in range(i + 1, n):
            if dist[i][j] == INF: continue #No connection, skip
            freq[dist[i][j]] += 1

    #If there are v pairs per distance (v >= 3), we can add v choose 3 into final answer
    ans = 0
    for v in freq.values():
        if v >= 3:
            ans += nCr(v,3)
    return ans

test_case = 1
for _ in range(test_case):
    n = 5
    A =  [[1, 2], [1, 3], [1, 4], [1, 5]]
    print(solve())