[Python] 1303번 전쟁 - 전투

✏️ 문제

 

문제 파악

입력으로 주어진 배열을 두 가지로 병사를 나눴다. 

• 1 → 'W'가 입력으로 들어온 위치, 즉 우리팀 병사
• 2 → 'B'가 입력으로 들어온 위치, 즉 상대팀 병사

병사를 구분하여 병사의 수를 세야하므로 카운트할 때 다음 위치에 있는 병사가 현재 병사와 같은 팀이라는 조건이 붙어야 한다.

 

또한 이 문제의 경우 가로, 세로로 이동해서 방문할 수 있으니 방향 배열을 다음과 같이 구성했다.

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

 

 

 

알고리즘

• 그래프 이론
• 그래프 탐색
• 너비 우선 탐색
• 깊이 우선 탐색

 

 

코드

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

m, n = map(int, input().split())
graph = [[0] * m for _ in range(n)]
visited = [[0] * m for _ in range(n)]
d = [(1, 0), (0, 1), (-1, 0), (0, -1)]
total = [0, 0]

for i in range(n):
  s = input()
  for j in range(m):
    if s[j] == 'W': # 우리팀
      graph[i][j] = 1
    else:
      graph[i][j] = 2

def bfs(x, y):
  q = deque([(x, y)])
  visited[x][y] = 1
  cnt = 0
  team = graph[x][y]
  
  while q:
    x, y = q.popleft()

    for i in range(4):
      dx = x + d[i][0]
      dy = y + d[i][1]

      if 0 <= dx < n and 0 <= dy < m and visited[dx][dy] == 0 and graph[dx][dy] == team: 
        cnt += 1
        visited[dx][dy] = 1
        q.append([dx, dy])
  
  return cnt + 1

for i in range(n):
  for j in range(m):
    if visited[i][j] == 0:
      if graph[i][j] == 1:
        total[0] += (bfs(i, j)**2)
      else:
        total[1] += (bfs(i, j)**2)

for i in range(2):
  print(total[i], end=' ')

 

 

---

 

⭐️ DFS와 BFS에 대해서 아직 잘 모른다면?

[알고리즘] DFS와 BFS