[코드트리 파이썬] 우리는 하나

문제

---

풀이

그래프를 탐색하는 문제이므로 BFS를 사용하였다.

 

n, k, u, d = map(int, input().split())
graph = [list(map(int, input().split())) for _ in range(n)]
visited = [[False] * n for _ in range(n)]
sumList = []

먼저 필요한 정보들을 입력받고 필요한 변수들을 선언하였다.

방문 여부를 체크해야하므로 visited배열을 사용하였다.

def BFS(i, j):
    q = deque()
    q.append((i, j))
    _sum = 1
    visited[i][j] = True

    while q:
        r, c = q.popleft()
        now = graph[r][c]
        for i in range(4):
            nr, nc = r + dr[i], c + dc[i]
            if 0 <= nr < n and 0 <= nc < n and u <= abs(now - graph[nr][nc]) <= d and not visited[nr][nc]:
                q.append((nr, nc))
                _sum += 1
                visited[nr][nc] = True
    
    sumList.append(_sum)

BFS는 다음과 같이 구현하였다. 

먼저 다음에 이동할 좌표를 저장할 q를 선언하고 시작 위치를 방문처리한다.

다음 q가 빌 때 까지 반복하는데, 현재 위치의 빌딩높이를 저장할 now변수를 할당해주고, for문을 이용하여 동서남북으로 이동한다.

이 때 이동할 위치가 격자 내부에 있으며, 이동할 위치와 현재위치간의 높이차가 u와d사이이며, 방문하지 않았을 경우 이동ㅎㄴ다.

또한 이동했으므로 이동한 개수를 저장하는 _sum변수에 1을 더해준다.

마지막으로 이동한 좌표를 방문처리한다.

while문이 종료 된 후 총 이동한 빌딩의 개수를 sumList 배열에 저장한다.

for i in range(n):
    for j in range(n):
        BFS(i, j)

sumList.sort(reverse = True)
res = 0
for i in range(k):
    res += sumList[i]
print(res)

BFS를 실행시키는 코드이다.

격자하나하나에 BFS코드를 실행시키면 모든 경우의수가 sumList의 저장된다.

이 때 우리가 원하는건 큰 수부터이므로 sumlist를 내림차순으로 정렬하고

원하는 개수만큼 res값을 더해 출력한다.

 

 

코드

from collections import deque

n, k, u, d = map(int, input().split())
graph = [list(map(int, input().split())) for _ in range(n)]
visited = [[False] * n for _ in range(n)]
sumList = []

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

def BFS(i, j):
    q = deque()
    q.append((i, j))
    _sum = 1
    visited[i][j] = True

    while q:
        r, c = q.popleft()
        now = graph[r][c]
        for i in range(4):
            nr, nc = r + dr[i], c + dc[i]
            if 0 <= nr < n and 0 <= nc < n and u <= abs(now - graph[nr][nc]) <= d and not visited[nr][nc]:
                q.append((nr, nc))
                _sum += 1
                visited[nr][nc] = True
    
    sumList.append(_sum)

for i in range(n):
    for j in range(n):
        BFS(i, j)

sumList.sort(reverse = True)
res = 0
for i in range(k):
    res += sumList[i]
print(res)