[Python] 백준 1303 전쟁 - 전투 - BFS

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

1303번: 전쟁 - 전투

문제

전쟁은 어느덧 전면전이 시작되었다. 결국 전투는 난전이 되었고, 우리 병사와 적국 병사가 섞여 싸우게 되었다. 그러나 당신의 병사들은 흰색 옷을 입고, 적국의 병사들은 파란색 옷을 입었기 때문에 서로가 적인지 아군인지는 구분할 수 있다. 문제는 같은 팀의 병사들은 모이면 모일수록 강해진다는 사실이다.

N명이 뭉쳐있을 때는 N2의 위력을 낼 수 있다. 과연 지금 난전의 상황에서는 누가 승리할 것인가? 단, 같은 팀의 병사들이 대각선으로만 인접한 경우는 뭉쳐 있다고 보지 않는다.

입력

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

출력

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

예제 입력 1 

5 5
WBWWW
WWWWW
BBBBB
BBBWW
WWWWW

예제 출력 1 

130 65

---

풀이

• 같은 팀의 병사들은 모이면 모일수록 강해지며 제곱의 위력 가짐
• 같은 팀 병사들이 대각선으로 인접할 경우 뭉쳐있다고 보지 않기 때문에, BFS를 사용하여 상하좌우 탐색 후 병사의 수 저장
• 위력의 합 출력 위해 위력 더하는 변수 2개(blue, white) 선언
• BFS가 끝날 때, 위력 반환하여 모두 더한 뒤 출력
• strip(): 병사 입력 받을 때 문자열 공백 삭제 함수 사용

코드 

from collections import deque;

# 상하좌우
dx = [0,0,1,-1]
dy = [1,-1,0,0]

# 병사의 수
cnt = 0

# 병사의 옷 색 (W, B)
color = ''

# 아군, 적군의 위력으로 제곱해서 더할 변수
blue = 0
white = 0

def bfs(x, y, color):
    # 현재 들어온 병사를 큐에 넣기 때문에 병사의 수(cnt)를 1로 시작
    # 방문 처리 위해 현재 좌표값을 True로 변경
    # 큐에 넣은 좌표값을 꺼내어 상하좌우 탐색 -> 해당 값이 있고 방문하지 않았다면, 큐에 넣고 병사의 수 증가
    # 병사 n명 뭉쳐있을 경우 제곱근의 위력 -> cnt^2 반환 
    cnt = 1
    q = deque()
    q.append((x,y))
    visited[x][y] = True

    while q:
        x, y = q.popleft()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]

            if nx < 0 or nx >= m or ny < 0 or ny >= n:
                continue
            if board[nx][ny] == color and not visited[nx][ny]:
                q.append((nx, ny))
                cnt += 1
                visited[nx][ny] = True  
    return cnt ** 2

if __name__ == "__main__":
    n, m = map(int, input().split())                        # 가로 세로 입력
    board = [list(input().strip()) for _ in range(m)]       # 병사의 옷 색 공백 없이 입력 받아 2차원 리스트 생성 
    visited = [[False] * n for _ in range(m)]               # 방문 확인 리스트

    for i in range(m):
        for j in range(n):
            if board[i][j] == 'W' and not visited[i][j]:    # 값이 'W'이고 방문하지 않았을 때 BFS 실행
                color = board[i][j]
                white += bfs(i, j, color)                   # 하얀 옷을 입은 병사의 위력 저장

            elif board[i][j] == 'B' and not visited[i][j]:  # 값이 'B'이고 방문하지 않았을 때 BFS 실행
                color = board[i][j]
                blue += bfs(i, j, color)                    # 파란 옷을 입은 병사의 위력 저장


print(white, blue)