문제

전쟁은 양 대륙의 군사들이 대결하는 것이다. 각 군사는 병사들로 이루어져 있다. 병사들은 같은 편 병사와 인접해 있을수록 강하다.

N명의 병사가 모여있을 때 위력은 N^2이다. 각 팀의 위력을 구하는 프로그램을 작성하시오.

풀이

BFS/DFS를 사용하여 연결된 병사들을 세는 문제이다.

코드

from collections import deque

N, M = map(int, input().split())
grid = [list(input().strip()) for _ in range(M)]
visited = set()

dxy = [(-1, 0), (1, 0), (0, -1), (0, 1)]

def bfs(x, y, team):
    q = deque([(x, y)])
    visited.add((x, y))
    count = 1
    
    while q:
        x, y = q.popleft()
        
        for dx, dy in dxy:
            nx, ny = x + dx, y + dy
            
            if not(0 <= nx < M and 0 <= ny < N):
                continue
            if (nx, ny) in visited:
                continue
            if grid[nx][ny] != team:
                continue
            visited.add((nx, ny))
            q.append((nx, ny))
            count += 1
    return count

white_power = 0
blue_power = 0

for i in range(M):
    for j in range(N):
        if (i, j) in visited:
            continue
        team = grid[i][j]
        count = bfs(i, j, team)
        if team == 'W':
            white_power += count ** 2
        else:
            blue_power += count ** 2

print(white_power, blue_power)

시간 복잡도

O(N × M)