F - Add One Edge 2 Editorial by evima

Another Solution

The objective is to count the number of vertex pairs \((u, v)\) (where \(u < v\)) that satisfy the following conditions:

  • Both vertices \(u\) and \(v\) have a degree of \(2\).
  • All vertices on the simple path connecting \(u\) and \(v\) (excluding \(u\) and \(v\) themselves) have a degree of \(3\).

This can be achieved by enumerating the connected components of the subgraph induced by all vertices with a degree of \(3\). For each connected component, let \(c\) be the number of degree \(2\) vertices adjacent to it. Then, the number of valid pairs from that connected component is \(c(c - 1)/2\).

Sample Implementation (Python):

from atcoder.dsu import DSU

N = int(input())
u, v = [0] * (N - 1), [0] * (N - 1)
deg = [0] * N
for i in range(N - 1):
    u[i], v[i] = map(int, input().split())
    u[i], v[i] = u[i] - 1, v[i] - 1
    deg[u[i]] += 1
    deg[v[i]] += 1

d = DSU(N)
c2 = [0] * N
for i in range(N - 1):
    if deg[u[i]] == deg[v[i]] == 3:
        d.merge(u[i], v[i])
    elif deg[u[i]] == 3 and deg[v[i]] == 2:
        c2[u[i]] += 1
    elif deg[u[i]] == 2 and deg[v[i]] == 3:
        c2[v[i]] += 1

ans = 0
for g in d.groups():
    c = 0
    for v in g:
        c += c2[v]
    ans += c * (c - 1) // 2
print(ans)

posted:
last update: