SWEA 2001번: 파리 퇴치 (D2) - Python 풀이 (feat. 달팽이의 저주)

문제 설명

SW Expert Academy

N x N 배열 안의 숫자는 해당 영역에 존재하는 파리의 개수를 의미한다.

아래는 N=5 의 예이다.

M x M 크기의 파리채를 한 번 내리쳐 최대한 많은 파리를 죽이고자 한다.

죽은 파리의 개수를 구하라!

예를 들어 M=2 일 경우 위 예제의 정답은 49마리가 된다.

[제약 사항]

1. N 은 5 이상 15 이하이다.
2. M은 2 이상 N 이하이다.
3. 각 영역의 파리 갯수는 30 이하 이다.

[입력]

가장 첫 줄에는 테스트 케이스의 개수 T가 주어지고, 그 아래로 각 테스트 케이스가 주어진다.
각 테스트 케이스의 첫 번째 줄에 N 과 M 이 주어지고, 다음 N 줄에 걸쳐 N x N 배열이 주어진다.

[출력]

출력의 각 줄은 '#t'로 시작하고, 공백을 한 칸 둔 다음 정답을 출력한다. (t는 테스트 케이스의 번호를 의미하며 1부터 시작한다.)

입력

10
5 2
1 3 3 6 7
8 13 9 12 8
4 16 11 12 6
2 4 1 23 2
9 13 4 7 3
6 3
29 21 26 9 5 8
21 19 8 0 21 19
9 24 2 11 4 24
19 29 1 0 21 19
10 29 6 18 4 3
29 11 15 3 3 29
...

input.txt

출력

#1 49
#2 159
...

output.txt

문제 요약) NN 숫자 배열에서 MM 범위에 있는 값들의 총합이 최댓값이 될 때, 최댓값을 출력해라.

풀이 과정

아니 이거 그냥 for문 4개로도 풀 수 있다고 하네요. 심지어 누적합 알고리즘이었습니다. 간단한 방법을 찾고 계시다면 다른 분 블로그 보시는 게 더 좋으실 것 같습니다. 최고의 방법은 아니지만 제가 생각한 로직을 정리해 보겠습니다.

저는 그렇게 풀지는 않았는데요. 좀 복잡하게 풀었는데 왜냐하면 이 문제를 보자마자 이틀 전에 푼 달팽이 숫자 문제가 생각나서 그 문제의 심화 버전인 줄 알고 응용해서 풀었습니다. 달팽이 문제는 우→하→좌→상 방향으로 움직이며 숫자를 1씩 증가해서 채우는 것이었다면, 이 문제는 2차원 배열 안에서 움직이면서 가변적인 MM 범위 내의 숫자의 합을 구해야 합니다. 가능한 모든 MM 범위의 숫자의 합을 구하고, 그 중 최댓값을 출력하는 방식으로 풀이했습니다.

그려지는 네모의 좌측 상단을 기준점으로 두고 모든 케이스를 확인해보니, rect_criteria가 NN의 좌표에서 (N-(M-1))(N-(M-1)) 크기만큼 움직인다는 것을 알 수 있었습니다.

그냥 풀면서 막 그린 거... 그림 이해에 방해될 것 같은 필기는 흰색으로 지웠는데 이해가 되실지 모르겠습니다

그리하여 달팽이 문제처럼 네모의 기준점을 (N-(M-1))(N-(M-1)) 범위 내에서 우→하→좌→상으로 움직이며 기준점으로부터 그려지는 MM 범위 안의 숫자 합을 answer 배열에 append 했습니다.

심지어 달팽이 문제랑 똑.같이 M*M ‘범위 밖으로 나가는 것 외에’ 해당 네모 이미 과거에 그려진 네모일 때에도 방향 전환을 해야 하기 때문에 rect_criteria라는 Bool 타입 2차원 배열을 만들어, 이미 기준점이 거쳐간 자리이면 True로 표시해 주었습니다. 여기까지 캐치하며 ‘와. 베리에이션 미쳤다. 완전 달팽이 심화 버전이네 ㅋㅋ.’ 라고 생각했습니다 ^^. 그 어디에도 우하좌상 방향으로 움직이라는 말도 없었는데… 그냥 2차원 배열 위에서 계속 네모를 그려가며 합을 구해야 하는 문제라서 똑같이 재귀로 풀면 되겠다고 생각했던 것 같습니다.

아무튼 그렇게 recursive라는 재귀함수를 구현했고, sum_bugs라는 함수로 좌측 상단 시작점 (x,y)에서부터 그려지는 M*M의 범위 안의 합을 구했습니다. 이때 2차원 배열 슬라이싱을 사용했습니다.

결론) 그려질 네모의 좌측 상단을 기준점 (x,y)으로 두고, 이 기준점을 달팽이 문제처럼 이동시키며 가능한 모든 M*M 범위 안의 값의 총합을 구하고 그 중 최댓값을 출력합니다.

정답 코드

# M * M 범위에 있는 값들의 총합이 최댓값이 될 때 찾기
 
T = int(input())
# 여러개의 테스트 케이스가 주어지므로, 각각을 처리합니다.
for test_case in range(1, T + 1):
    N, M = map(int, input().split())
    bugs = [list(map(int, input().split())) for _ in range(N)]
    # M*M의 네모의 좌측 상단의 기준점이 지나갈 수 있는 범위 (N-(M-1))*(N-(M-1))
    # 에서 기준점이 해당 위치를 지나갔는지 여부를 저장하기 위한 2차원 배열
    rect_criteria = [[False] * (N-(M-1)) for _ in range(N-(M-1))]
 
    # 가능한 모든 M*M 네모 안의 총합을 담을 1차원 배열
    answer = []
 
    # 우->하->좌->상 으로의 방향 전환을 위한 변수 (달팽이 문제의 관성 때문임 꼭 이 방향 아니어도 됩니다...)
    dx = [0, 1, 0, -1]
    dy = [1, 0, -1, 0]
 
    # 기준점 (x,y)를 좌측 상단으로 두는 M*M의 네모를 그리고
    # 그 안의 값의 총합을 구해 answer에 추가해주는 함수
    def sum_bugs(x, y):
        value = 0
        # 네모가 그려질 범위를 슬라이싱 (<- 행을 슬라이싱)
        ranged_bugs = bugs[x:x+M]
        # 열을 슬라이싱
        for i in range(M):
            ranged_bugs[i] = ranged_bugs[i][y:y+M]
        # 범위 내 값의 총합
        for bug in ranged_bugs:
            value += sum(bug)
        # answer에 추가
        answer.append(value)
        
    # 가능한 모든 M*M의 범위를 찾는 재귀함수
    def recursive(x, y, d):
        # 영역 표시. 현재 기준점 (x,y) 왔다 갑니다 ~
        rect_criteria[x][y] = True
        # 현재 기준점에서부터 그려지는 M*M 네모 안의 값의 총합을 구합니다
        sum_bugs(x, y)
         
        # x,y 좌표 이동
        nx = x + dx[d]
        ny = y + dy[d]
 
        # 2차원이라 1차원으로 flat 시키고
        flat_list = sum(rect_criteria, [])
        # 모두 True 값을 가질 때 (= 가능한 모든 기준점을 확인했을 때) 재귀 종료
        if all(flat_list):
            return
        # 이동하려는 좌표가 가능한 범위 밖이거나
        # 이미 과거에 지나갔던 기준점이라면
        if nx < 0 or ny < 0 or nx >= N-(M-1) or ny >= N-(M-1) or rect_criteria[nx][ny]:
            # 방향 전환
            d = (d+1) % 4
            # 처음 x,y 값 그대로 재귀 호출
            recursive(x, y, d)
        # 이동하려는 좌표가 가능한 범위이고 처음 지나가는 기준점이라면
        else:
            # 이동시킨 좌표 값으로 재귀 호출
            recursive(nx, ny, d)
		
    # 기준점은 (0,0)에서, 방향은 오른쪽으로 재귀 시작
    recursive(0, 0, 0)
    print("#%d %d" % (test_case, max(answer)))

이상입니다.

가 아니고 양심상 모범 답안도 정리를 하겠습니다.

1. 무지성 for문 4개 쓰는 방법. (N이 작기 때문에 이렇게 해도 문제 X 아마 그래서 D2인 것으로 추정)

# 테스트 케이스 개수
T = int(input())
for test_case in range(1, T + 1):
    N, M = map(int, input().split())
    arr = [list(map(int,input().split())) for _ in range(N)]

    # 최댓값
    max_v = 0
    # (0, 0)에서 (N-M, N-M)까지 확인
    for i in range(N-M+1):
        for j in range(N-M+1):
            arr_sum = 0
            # M*M 사각형 안의 값의 총합 구하기
            for x in range(i,i+M):
                for y in range(j,j+M):
                    arr_sum += arr[x][y]
            # 구한 총합이 현재 최댓값보다 크면 최댓값 갱신
            max_v = max(max_v, arr_sum)
    print(f'#{test_case} {max_v}')

2. 누적합 알고리즘을 적용한 풀이

T = int(input())

for test_case in range(1, T + 1):
    N, M = map(int, input().split())
    bugs = [list(map(int, input().split())) for _ in range(N)]

    # 누적합을 저장할 배열 생성
    cumulative_sum = [[0] * (N + 1) for _ in range(N + 1)]
    # 인덱스 1에서 N까지 확인
    for i in range(1, N + 1):
        for j in range(1, N + 1):
            # (1,1)부터 (i,j)까지의 네모 영역 합
            cumulative_sum[i][j] = (
                cumulative_sum[i - 1][j] # 위쪽 누적합
                + cumulative_sum[i][j - 1] # 왼쪽 누적합
                - cumulative_sum[i - 1][j - 1] # 겹치는 부분 제거
                + bugs[i - 1][j - 1] # 현재 위치 값 추가
            )

    # 최대 파리 수 계산
    answer = 0
    for i in range(N - M + 1):
        for j in range(N - M + 1):
            # (i, j)부터 (i+M-1, j+M-1)까지의 네모 영역 합을 누적합으로 구하는 공식
            total = (
                cumulative_sum[i + M][j + M] # 전체 영역 누적합
                - cumulative_sum[i][j + M] # 위쪽 불필요한 영역 제거
                - cumulative_sum[i + M][j] # 왼쪽 불필요한 영역 제거
                + cumulative_sum[i][j] # 위에서 2번 제거된 겹치는 영역 복구
            )
            # 최댓값을 answer에 저장
            answer = max(answer, total)

    print(f"#{test_case} {answer}")

출처: https://nahwasa.com/entry/누적-합prefix-sum-2차원-누적합prefix-sum-of-matrix-with-java#2차원_누적_합_(prefix_sum_of_matrix)

하단 이중 for문에서 total을 구하는 공식은 빨간 영역의 누적합을 구하기 위해 전체 영역 누적합 - (초록 영역 누적합) - (노란 영역 누적합) + (노랑 초록 겹친 부분 누적합 (복구))를 한 것입니다.

어우 2차원 누적합 어렵습니다.

이상입니다.