[SWEA/D2] 2001. 파리 퇴치

[문제 링크]

https://swexpertacademy.com/main/code/problem/problemDetail.do?problemLevel=2&contestProbId=AV5PzOCKAigDFAUq&categoryId=AV5PzOCKAigDFAUq&categoryType=CODE&problemTitle=&orderBy=FIRST_REG_DATETIME&selectCodeLang=PYTHON&select-1=2&pageSize=10&pageIndex=1

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부터 시작한다.)

---

제출 코드

t = int(input())

# 누적합 계산 함수 생성
def prefix_sum(n, arr):
    prefix = [[0 for _ in range(n+1)] for _ in range(n+1)]
    for i in range(1, n+1):
        for j in range(1, n+1):
            prefix[i][j] = prefix[i][j-1] + prefix[i-1][j] + arr[i-1][j-1] - prefix[i-1][j-1]
    return prefix

for i in range(1, t+1):
    n, m = map(int, input().split())
    arr = [list(map(int, input().split())) for _ in range(n)]
    prefix = prefix_sum(n, arr)
    
    # 구간 합 계산하기
    m_sum = []
    for x in range(m,n+1):
        for y in range(m, n+1):
            flies = prefix[x][y] - prefix[x-m][y] - prefix[x][y-m] + prefix[x-m][y-m]
            m_sum.append(flies)
    print('#'+str(i), max(m_sum))

 

풀이

• 누적합 알고리즘을 통해 개선된 시간복잡도로 누적합을 계산한다.
• m X m 크기의 구간합을 계산한 후, max_flies와 비교해 큰 값을 저장한다.

---

알아둘 사항

누적합 알고리즘

• 누적합(Prefix Sum)은 배열이나 행렬에서 특정 구간의 합을 빠르게 구할 수 있도록 사전 계산을 사용하는 알고리즘.
• 특정 구간의 합을 자주 계산해야 할 때 매우 유용함.
• 반복적인 합산 과정을 줄이고, 특정 구간의 합을 효율적으로 계산 가능.
  • 초기 O(n) 또는 O(n^2) 연산을 통해 구간 합을 O(1)로 계산 가능

2차원 누적합을 계산할 때에는, 점화식 구성시 편의를 위해 첫 행과 첫 열을 0으로 구성한다.

아래에 참고한 글을 읽어보면 도움이 될 듯.

구간합을 구하기 위해서는, 누적합에서 원하는 위치 이전까지의 누적합을 빼주면 된다.

 

아래 글을 참고해서 이해했다.

https://code-angie.tistory.com/22

[알고리즘 / Python] 누적합 (Prefix Sum)

 

• 보다 개선된 코드

t = int(input())

def prefix_sum(n, arr):
    prefix = [[0] * (n+1) for _ in range(n+1)]
    for i in range(1, n+1):
        for j in range(1, n+1):
            prefix[i][j] = prefix[i][j-1] + prefix[i-1][j] + arr[i-1][j-1] - prefix[i-1][j-1]
    return prefix

for i in range(1, t+1):
    n, m = map(int, input().split())
    arr = [list(map(int, input().split())) for _ in range(n)]
    prefix = prefix_sum(n, arr)
    
    max_flies = 0
    for x in range(m,n+1):
        for y in range(m, n+1):
            flies = prefix[x][y] - prefix[x-m][y] - prefix[x][y-m] + prefix[x-m][y-m]
            max_flies = max(max_flies, flies)
    print(f"#{i} {max_flies})

 

m_sum으로 죽은 파리수를 모두 저장해서 max로 계산하는 것보다, max_flies로 그때그때 비교해서 저장하는 것이 메모리 효율이 더 좋다.