Submission #1213340

#	Time	Username	Problem	Language	Result	Execution time	Memory
1213340		spike1236	Duathlon (APIO18_duathlon)	Pypy 3	0 / 100	1120 ms	803720 KiB

count_triplets

import sys

sys.setrecursionlimit(3000000)
input = sys.stdin.readline

n, m = map(int, input().split())
adj = [[] for _ in range(n + 1)]
for _ in range(m):
    u, v = map(int, input().split())
    adj[u].append(v)
    adj[v].append(u)

disc = [0] * (n + 1)
low = [0] * (n + 1)
is_art = [False] * (n + 1)
time_dfs = 0

edge_stack = []
bcc_list = []
vertex_to_bcc = [[] for _ in range(n + 1)]


def dfs_tarjan(u, parent):
    global time_dfs
    disc[u] = low[u] = time_dfs + 1
    time_dfs += 1
    child_count = 0

    for v in adj[u]:
        if disc[v] == 0:
            child_count += 1
            edge_stack.append((u, v))
            dfs_tarjan(v, u)
            low[u] = min(low[u], low[v])

            if (parent != -1 and low[v] >= disc[u]) or (
                parent == -1 and child_count > 1
            ):
                is_art[u] = True

            if low[v] >= disc[u]:
                new_vertices = set()
                while True:
                    x, y = edge_stack.pop()
                    new_vertices.add(x)
                    new_vertices.add(y)
                    if (x, y) == (u, v) or (x, y) == (v, u):
                        break

                bcc_id = len(bcc_list)
                bcc_list.append(list(new_vertices))
                for w in new_vertices:
                    vertex_to_bcc[w].append(bcc_id)

        elif v != parent and disc[v] < disc[u]:
            low[u] = min(low[u], disc[v])
            edge_stack.append((u, v))


for i in range(1, n + 1):
    if disc[i] == 0:
        dfs_tarjan(i, -1)

        if edge_stack:
            new_vertices = set()
            while edge_stack:
                x, y = edge_stack.pop()
                new_vertices.add(x)
                new_vertices.add(y)
            bcc_id = len(bcc_list)
            bcc_list.append(list(new_vertices))
            for w in new_vertices:
                vertex_to_bcc[w].append(bcc_id)

num_bcc = len(bcc_list)

totalNodes = n + num_bcc
bct_adj = [[] for _ in range(totalNodes + 1)]
B_sz = [0] * (num_bcc + 1)
for bcc_id, verts in enumerate(bcc_list, start=0):
    node_id = n + 1 + bcc_id
    B_sz[bcc_id] = len(verts)
    art_count = 0
    for v in verts:
        if is_art[v]:
            art_count += 1
            bct_adj[node_id].append(v)
            bct_adj[v].append(node_id)

weight = [0] * (totalNodes + 1)
for v in range(1, n + 1):
    if is_art[v]:
        weight[v] = 1

for bcc_id in range(num_bcc):
    node_id = n + 1 + bcc_id

    art_in_block = 0
    for v in bcc_list[bcc_id]:
        if is_art[v]:
            art_in_block += 1
    weight[node_id] = B_sz[bcc_id] - art_in_block

parent = [-1] * (totalNodes + 1)
size_subtree = [0] * (totalNodes + 1)
comp_total = [0] * (totalNodes + 1)
visited_bct = [False] * (totalNodes + 1)

for start in range(1, totalNodes + 1):
    if (start <= n and not is_art[start]) or visited_bct[start]:
        continue

    stack = [(start, -1, 0)]
    visited_bct[start] = True
    parent[start] = 0
    comp_nodes = []
    while stack:
        node, par, idx = stack[-1]
        if idx == 0:
            comp_nodes.append(node)
        if idx < len(bct_adj[node]):
            nei = bct_adj[node][idx]
            stack[-1] = (node, par, idx + 1)
            if not visited_bct[nei]:
                visited_bct[nei] = True
                parent[nei] = node
                stack.append((nei, node, 0))

        else:
            total_w = weight[node]
            for w in bct_adj[node]:
                if parent[w] == node:
                    total_w += size_subtree[w]
            size_subtree[node] = total_w
            stack.pop()

    totW = size_subtree[start]
    for u in comp_nodes:
        comp_total[u] = totW

answer = 0

for c in range(1, n + 1):
    if not is_art[c]:
        bcc_id = vertex_to_bcc[c][0]
        D = B_sz[bcc_id] - 1
        if D >= 2:
            answer += D * (D - 1)

    else:
        sum_S = 0
        sum_S_sq = 0
        sum_intra = 0

        for bnode in bct_adj[c]:
            if parent[bnode] == c:
                S_j = size_subtree[bnode]
            else:
                S_j = comp_total[c] - size_subtree[c]

            sum_S += S_j
            sum_S_sq += S_j * S_j

            bcc_id = bnode - n - 1
            D_j = B_sz[bcc_id] - 1

            L_j = 0
            sum_l2 = 0
            for a in bct_adj[bnode]:
                if a == c:
                    continue

                if parent[a] == bnode:
                    comp_ab = size_subtree[a]
                else:
                    comp_ab = comp_total[bnode] - size_subtree[bnode]
                l_i = comp_ab - 1
                L_j += l_i
                sum_l2 += l_i * l_i

            intra = 0
            if D_j >= 2:
                intra += D_j * (D_j - 1)
            intra += L_j * L_j - sum_l2
            intra += 2 * (D_j - 1) * L_j

            sum_intra += intra

        cross_pairs = sum_S * sum_S - sum_S_sq
        answer += cross_pairs + sum_intra

print(answer)

Compilation message (stdout)

Compiling 'count_triplets.py'...

=======
  adding: __main__.pyc (deflated 43%)

=======

• Subtask 1: 0 / 5
• Subtask 2: 0 / 11
• Subtask 3: 0 / 8
• Subtask 4: 0 / 10
• Subtask 5: 0 / 13
• Subtask 6: 0 / 15
• Subtask 7: 0 / 20
• Subtask 8: 0 / 8
• Subtask 9: 0 / 10