[백준] 4195번: 친구 네트워크 (retry)

https://www.acmicpc.net/problem/4195

initial

I tried solving via DFS but I had some troubles

1) how to initalise visited list. We are not going by the integer index but the string of the names so I asked chatgpt and she recommended using set()

2) how to start off this dfs

3) how to get the number of friends

wrong code: not sure why

from collections import defaultdict

def dfs(graph, visited):
    count = 1
    for i in range(len(graph)):
        for j in graph[i]:
            if not visited[j]:
                visited[j] = True
                count += dfs(graph, visited)
    return count

n = int(input())
graph = defaultdict(list)

for _ in range(n):
    m = int(input())
    for _ in range(m):
        a, b = map(str, input().split())
        graph[a].append(b)
        graph[b].append(a)

visited = set()
visited.add(list(graph.keys())[0])

ans = dfs(graph, visited)
print(ans)

solution

https://chancoding.tistory.com/50

I also thought of union find but not sure how to update the numbers of friends. We can use a separate dictionary to store the number of friends. We create a number dict and store the answer on the parent key node like {fred:2} by adding the number of incoming child, which is gonna be initialised as 1 if we have not seen this node before (else we just take the stored value of that node)

def find(a):
    if parents[a] != a:
        parents[a] = find(parents[a])
    return parents[a]


def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parents[b] = a
        counter[a] += counter[b]


for _ in range(int(input())):
    parents = {}
    counter = {}

    for i in range(int(input())):
        a, b = input().split()

        if a not in parents:
            parents[a] = a
            counter[a] = 1
        if b not in parents:
            parents[b] = b
            counter[b] = 1

        union(a, b)
        print(counter[find(a)])

very important that in the union, we get the parent and update the counter of the parent, not the child nodes.