[BOJ] 11660 : 구간 합 구하기 5 [누적합](Python)

문제 링크

https://www.acmicpc.net/problem/11660

11660번: 구간 합 구하기 5

소스 코드

import sys

def main():
    n, m = map(int, input().split())

    graph = [[0] * (n+1)]
    for _ in range(n):
        nums = [0] + list(map(int, sys.stdin.readline().strip().split()))
        graph.append(nums)
    
    for row in range(1,n+1):
        for col in range(2,n+1):
            graph[row][col] += graph[row][col-1] # 열 누적
    
    for row in range(2,n+1): # 행 누적
        for col in range(1,n+1):
            graph[row][col] += graph[row-1][col]

    for _ in range(m): # 누적합 계산(중복은 덧셈)
        x1, y1, x2, y2 = map(int, sys.stdin.readline().strip().split())
        print(graph[x2][y2] - graph[x1-1][y2] - graph[x2][y1-1] + graph[x1-1][y1-1])

if __name__ == '__main__':
    main()

시간초과 풀이

import sys

def main():
    n, m = map(int, input().split())

    graph = [[] for _ in range(n)]
    for i in range(n):
        temp = list(map(int, sys.stdin.readline().strip().split()))
        graph[i] = temp
        for j in range(1,n):
            graph[i][j] = graph[i][j-1] + graph[i][j]

    questions = []
    for _ in range(m):
        x1, y1, x2, y2 = map(int, sys.stdin.readline().strip().split())
        questions.append((x1-1, x2, y1-2, y2-1))

    for x1, x2, y1, y2 in questions:
        prefix(graph, x1, x2, y1, y2)

def prefix(graph, x1, x2, y1, y2):
    answer = 0

    if y1 == -1:
        for row in range(x1, x2):
            answer += graph[row][y2]
    else:
        for row in range(x1, x2):
            answer += graph[row][y2] - graph[row][y1]
    
    print(answer)


if __name__ == '__main__':
    main()

 

해설

• 다이나믹 프로그래밍, 누적합

시간초과를 받은 코드를 작성할 때도 나름 신경을 썼는데 확실히 깔끔하게 네 번만 계산을 하는 것과는 시간 차이가 많이 나네요.

 

시간복잡도

사실 두 풀이 모두 시간복잡도는 O(N^2)입니다.

어차피 2차원 리스트에 대해 모든 행과 열을 탐색하기 때문이죠.

하지만 시간 초과를 받은 풀이가 각 행의 누적합을 새로 계산(x2-x1+1번)하는 것과 달리 정석 풀이는 딱 네 번만 계산하면 끝납니다.

 

즉 탐색과는 별개로, 정답을 내기 위해 요구되는 계산의 시간 복잡도가 각각 최대 O(N)과 O(1)이기 때문에 시간 차이가 발생합니다.

M은 최대 100,000의 값을 가질 수 있으므로 만약 N이 1,000에 가까운 숫자가 된다면 시간 초과가 발생할 가능성이 높습니다.

 

그래프 초기화

누적합 개념을 이용하기 위해서 그래프를 잘 초기화해야 합니다.

특히나 마지막에 주어지는 x1, y1, x2, y2 의 값을 그대로 사용하기 위해서 그래프의 맨 앞에 0을 추가해서 초기화해줍니다.

graph = [[0] * (n+1)]
for _ in range(n):
    nums = [0] + list(map(int, sys.stdin.readline().strip().split()))
    graph.append(nums)

 

그리고 이를 각각 행별로, 열별로 누적시켜줍니다.

그 이유에 대해서는 뒤에 그림과 함께 설명하겠습니다.

for row in range(1,n+1):
    for col in range(2,n+1):
        graph[row][col] += graph[row][col-1] # 열 누적

for row in range(2,n+1): # 행 누적
    for col in range(1,n+1):
        graph[row][col] += graph[row-1][col]

 

누적합

문제의 예시를 자세히 살펴보도록 하죠.

왼쪽의 그래프를 오른쪽으로 압축시킵니다.

이는 각 행, 각 열의 값을 누적시킨 것입니다.

 

예를 들어 (2, 2) 좌표의 값 8은 (1, 1) 부터 (2, 2) 까지 이르는 모든 값을 더한 것입니다.

이게 곧 누적합이고 우리는 이 누적합의 차를 이용해서 원하는 값을 빠르게 구할 수 있습니다.

 

예를 들어 (3, 4)에 해당하는 값은 42입니다.

우리는 (2, 2) 부터 (3, 4) 까지의 값을 구하고 싶은 것이므로,

'빨강' - '파랑' - '초록' + '보라' 를 구해주면 정답이 됩니다.

 

(3, 4)에 저장된 값은 (2,2) 바깥의 값들도 포함하고 있는 것이기 때문에 이를 제외하기 위해 '파랑'과 '초록'을 빼주는 것이고, 중복으로 빠진 값을 '보라'로 더해주는 것입니다.

따라서 위 예시의 인덱스를 주의 깊게 확인해보고 코드로 구현하면 다음과 같습니다.

for _ in range(m): # 누적합 계산(중복은 덧셈)
    x1, y1, x2, y2 = map(int, sys.stdin.readline().strip().split())
    print(graph[x2][y2] - graph[x1-1][y2] - graph[x2][y1-1] + graph[x1-1][y1-1])

 

요약하면 각 행과 열에 담긴 값들을 누적시켜 그래프를 구성하고, 누적합의 합차를 이용하면 아주 빠르게 계산이 가능하다! 는 것입니다.