백준 1303 - 전쟁 - 전투 (파이썬)

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

풀이

이 문제는 BFS로 해결이 가능하다.

주어진 전장 지도에서 각 팀의 전투력을 계산하는 문제이며, 전투력은 같은 팀의 병사가 상하좌우로 연결된 병사들의 수의 제곱으로 계산된다.

• 방문 여부를 저장하는 visit 리스트
  • 같은 병사가 중복 탐색되지 않도록 하기 위해 사용
• BFS를 활용하여 같은 색상의 병사를 탐색
  • queue에 시작 병사의 좌표를 넣고 탐색
  • 같은 색상이면서 방문하지 않은 병사들을 queue에 추가하면서 탐색을 확장
  • 탐색이 끝나면 해당 병사의 그룹의 크기를 count * count로 계산하여 result에 추가

시간 복잡도 분석

• 각 병사를 한 번만 방문하므로
• O(N * M), 즉 주어진 전장의 크기만큼 수행
• BFS로 탐색하여 같은 팀 병사들을 묶으므로 효율적

코드

# 백준 1303 - 전쟁 - 전투
# 분류 : DFS/BFS

from collections import deque

N, M = map(int, input().split())

A = [""] * M

for i in range(M) :
    A[i] = input()

def solve(color) :
    result = 0
    visit = [[False] * N for _ in range(M)]
    queue = deque()

    dr = [1, -1, 0, 0]
    dc = [0, 0, 1, -1]

    for i in range(M) :
        for j in range(N) :
            if A[i][j] != color :
                continue

            if not visit[i][j] :
                queue.append((i, j))
                visit[i][j] = True
                count = 1

                while len(queue) != 0 :
                    r, c = queue.popleft()

                    for k in range(4) :
                        nr, nc = r + dr[k], c + dc[k]
                        if nr < 0 or M <= nr or nc < 0 or N <= nc :
                            continue
                        
                        if A[nr][nc] != color :
                            continue

                        if not visit[nr][nc] :
                            queue.append((nr, nc))
                            visit[nr][nc] = True
                            count += 1
                
                result += count * count
    return result

print(solve('W'), solve('B'))