백준 1303: 전쟁 - 전투 Python (BFS)

 

코드의 중복을 막기 위해서, bfs 함수의 인자로 현재 위치 뿐만 아니라 색도 함께 전달했다.

방문하지 않았으면서 색이 일치하는 것들만 counting 하고, 마지막에 제곱해서 결괏값에 더한다.

from collections import deque
import sys
input = sys.stdin.readline


def bfs(si, sj, color):
    army = 0
    q = deque()
    q.append((si, sj))
    visited[si][sj] = 1

    while q:
        si, sj = q.popleft()
        army += 1
        for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            ni, nj = si+di, sj+dj
            if 0 <= ni < M and 0 <= nj < N and not visited[ni][nj]:
                if arr[ni][nj] == color:
                    q.append((ni, nj))
                    visited[ni][nj] = 1
    return army


N, M = map(int, input().split())    # N: 가로(j), M: 세로(i)
arr = [list(input().rstrip()) for _ in range(M)]
visited = [[0] * N for _ in range(M)]   # 방문체크

result_w = 0
result_b = 0
for i in range(M):
    for j in range(N):
        if not visited[i][j]:
            if arr[i][j] == 'W':
                result_w += (bfs(i, j, 'W') ** 2)
            else:
                result_b += (bfs(i, j, 'B') ** 2)
print(result_w, result_b)