백준 - 단지번호붙이기

• DFS, 스택 이용
• 테스트케이스는 맞았는데 틀렸다고 표시된 풀이
  * 틀린 이유: 아파트의 개수를 카운트하는 cnt 변수를 0으로 시작하게 되면 아파트가 하나밖에 없는 단지의 경우 0으로 카운트 되어짐.
  

과정

1. (0, 0)부터 시작해서 아파트가 존재하는지 판단한다.
2. 현재 노드에 아파트가 존재한다면, 4 방향으로 아파트가 존재하는지 확인(dfs 함수 호출)한다.
3. 방문한 아파트는 0으로 만들어 준다.
4. 연결된 그래프의 개수를 카운트

틀린 풀이

from sys import stdin
n = int(input())
matrix = [stdin.readline().rstrip() for _ in range(n)]

for i in range(n):
    matrix[i] = [int(j) for j in matrix[i]]

def dfs(i, j): # stack
    stack = [(i, j)]
    dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]
    cnt = 0

    while stack:
        a, b = stack.pop()

        for i in range(4):
            nx = a + dx[i]
            ny = b + dy[i]

            if 0 <= nx < n and 0 <= ny < n and matrix[nx][ny] == 1:
                matrix[nx][ny] = 0
                cnt += 1
                stack.append((nx, ny))
    return cnt

entire_sum = 0
each_sum = []
for i in range(n):
    for j in range(n):
        if matrix[i][j] == 1:
            each_sum.append(dfs(i, j))
            entire_sum += 1

print(entire_sum)
for item in sorted(each_sum):
    print(item)

다시 푼 것

from sys import stdin
n = int(input())
matrix = [stdin.readline().rstrip() for _ in range(n)]

for i in range(n):
    matrix[i] = [int(j) for j in matrix[i]]

def dfs(i, j): # stack
    stack = [(i, j)]
    dx, dy = [-1, 1, 0, 0], [0, 0, -1, 1]
    
    cnt = 1 # cnt를 0으로 시작하면, 집이 하나밖에 없는 단지의 집이 카운트되지 않음.
    matrix[i][j] = 0 # 방문 표시

    while stack:
        a, b = stack.pop()

        for i in range(4):
            nx = a + dx[i]
            ny = b + dy[i]

            if 0 <= nx < n and 0 <= ny < n and matrix[nx][ny] == 1:
                matrix[nx][ny] = 0
                cnt += 1
                stack.append((nx, ny))
    return cnt

entire_sum = 0
each_sum = []

for i in range(n):
    for j in range(n):
        if matrix[i][j] == 1:
            each_sum.append(dfs(i, j))
            entire_sum += 1

print(entire_sum)
for item in sorted(each_sum):
    print(item)