ABC 157 D --Solve Friend Suggestions in Python!

The recent D problem is often unsolvable due to its bones, and I feel sad nowadays. I was wondering how to solve it even after it was over, so I tried my best to solve it.

The keyword this time is __ "undirected graph" __ "connected component" __. The points are __ "DFS" __ and __ "set" __!

Link to problem ABC 157 D --Friend Suggestions

Answers during the contest

I noticed that it was an undirected graph and immediately decided to solve it with DFS. From each vertex, DFS is used to count points from connected components that are neither friends nor blocks. If I thought that the double loop would definitely not be in time, it was TLE as expected. This is the time up during the contest.

from collections import deque

N,M,K = map(int,input().split())
friend = [list(map(int,input().split())) for _ in range(M)]
block = [list(map(int,input().split())) for _ in range(K)]

g = [[False]*(N+1) for _ in range(N+1)]
for i in friend:
    g[i[0]][i[1]] = True
    g[i[1]][i[0]] = True


b = [[False]*(N+1) for _ in range(N+1)]
for i in block:
    b[i[0]][i[1]] = True
    b[i[1]][i[0]] = True
    
stack=deque()
ans = []

for i in range(1,N+1):
    cnt = 0
    visited = [0] * (N+1)
    visited[i] = 1 
    stack.append(i)
    while stack:
        n = stack.pop()
        for j in range(1,N+1):
            if g[n][j] and visited[j]==0:
                stack.append(j)
                visited[j] = 1
                if g[i][j] == False and b[i][j]==False:
                    cnt+=1
    ans.append(cnt)
print(*ans)

Answers after the contest

When using in, I heard that Set is faster than List, so I decided to use it. As a result of fighting TLE through various trials and errors, I received AC below.

Accumulate the connected components in the link and find the linked components. After searching for connected components, It is calculated by subtracting the number of friend relationships and block relationships to be drawn for each point of the connected component.

from collections import deque

N,M,K = map(int,input().split())
friend = [list(map(int,input().split())) for _ in range(M)]
block = [list(map(int,input().split())) for _ in range(K)]

f = [set() for _ in range(N+1)]
b = [set() for _ in range(N+1)]

for i,j in friend:
    f[i].add(j)
    f[j].add(i)
for i,j in block:
    b[i].add(j)
    b[j].add(i)

stack=deque()
ans = [0]*(N+1)
visited = [0]*(N+1)

for i in range(1,N+1):
    if visited[i]:
        continue
    link = {i}
    visited[i] = 1 
    stack.append(i)
    while stack:
        n = stack.pop()
        for j in f[n]:
            if visited[j]==0:
                stack.append(j)
                visited[j] = 1
                link.add(j)
    for i in link:
        ans[i] = len(link) - len(link & f[i]) - len(link & b[i]) - 1
print(*ans[1:])

Ingenuity point

Since it is obviously useless to check the searched, I added not to search the searched.

if visited[i]:
    continue

I was checking if all the vertices were friends for a certain point, I changed to search only for those who have a friendship with that point.

for j in range(1,N+1):
# ↓↓↓
for j in f[n]:

I devised a way to subtract the number of friendships or block relations from the size of the connected component.

ans[i] = len(link) - len(link & f[i]) - len(link & b[i]) - 1

In other words, what does that mean ... By using the intersection, we are finding the number of friendships and block relations from the connected components.

f = [set for _ in range(10)]
b = [set for _ in range(10)]

#Connected components
link = {1,2,4,9,10}
#1 and 4/10 are friends
f[1] = {4,10}
#1 is blocking 5/6/9
b[1] = {5,6,9}

print(link&f[1], link&b[1])
print(len(link&f[1]), len(link&b[1]))

{10, 4} {9}
2 1