백준 1303번. 전쟁 - 전투 (Python / 파이썬)

문제

전쟁은 어느덧 전면전이 시작되었다. 결국 전투는 난전이 되었고, 우리 병사와 적국 병사가 섞여 싸우게 되었다. 그러나 당신의 병사들은 흰색 옷을 입고, 적국의 병사들은 파란색 옷을 입었기 때문에 서로가 적인지 아군인지는 구분할 수 있다. 문제는 같은 팀의 병사들은 모이면 모일수록 강해진다는 사실이다.
N명이 뭉쳐있을 때는 N2의 위력을 낼 수 있다. 과연 지금 난전의 상황에서는 누가 승리할 것인가? 단, 같은 팀의 병사들이 대각선으로만 인접한 경우는 뭉쳐 있다고 보지 않는다.

입력

첫째 줄에는 전쟁터의 가로 크기 N, 세로 크기 M(1 ≤ N, M ≤ 100)이 주어진다. 그 다음 두 번째 줄에서 M+1번째 줄에는 각각 (X, Y)에 있는 병사들의 옷색이 띄어쓰기 없이 주어진다. 모든 자리에는 병사가 한 명 있다. B는 파란색, W는 흰색이다. 당신의 병사와 적국의 병사는 한 명 이상 존재한다.

출력

첫 번째 줄에 당신의 병사의 위력의 합과 적국의 병사의 위력의 합을 출력한다.

---

문제 풀이 

적과 아군은 N명이 뭉쳐있을 때 N 제곱의 위력을 낼 수 있다. 

따라서 전쟁터를 탐색하면서 적과 아군이 뭉쳐있는 구간을 찾아내야 하며, 몇 명이 뭉쳐있는지 알아야 한다. 

적이 뭉쳐 있는 공간, 아군이 뭉쳐 있는 공간을 따로 계산해야 하기 때문에 옷색 정보가 중요하다. 

위의 그림처럼 아군이 뭉쳐있는 구간은 노란색 구간으로 총 2개이며 각각의 구간은 

9명, 7명으로 구성되어있다. 따라서 아군의 위력은 9의 제곱 + 7의 제곱 = 130이다. 

마찬가지로 적군이 뭉쳐있는 구간은 연두색 구간으로 총 2개이며 각각의 구간은 

1명, 8명으로 구성되어있다. 따라서 적군의 위력은 1의 제곱 + 8의 제곱 = 65이다. 

 

이처럼 적과 아군이 각각 뭉쳐있는 구간과 각 구간의 병사 수를 탐색하기 위해 BFS를 활용하였다. 

적의 옷색일 때, 그리고 아군의 옷색일 때 두 가지 경우로 나누어 BFS를 활용해 탐색하면 

적과 아군이 뭉쳐있는 구간 각각의 병사 수를 반환해준다. 

따라서 반환된 값을 제곱하여 다 더해주면 원하는 답을 출력할 수 있다. 

---

My Code

from collections import deque 

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

def bfs(x, y, color):
    q = deque() 
    q.append((x, y))
    visited[x][y] = True
    count = 1
    while q:
        x, y = q.popleft() 
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < m:
                if not visited[nx][ny] and array[nx][ny] == color:
                    visited[nx][ny] = True 
                    q.append((nx ,ny))
                    count += 1 
    return count 

m, n = map(int,input().split())

array = []
for _ in range(n):
    array.append(list(input()))

visited = [[False] * m for _ in range(n)]

blue = 0
white = 0
for i in range(n):
    for j in range(m):
        if not visited[i][j] and array[i][j] == 'B':
            blue += bfs(i, j, array[i][j]) ** 2
        elif not visited[i][j] and array[i][j] == 'W':
            white += bfs(i, j, array[i][j]) ** 2

print(white, blue)