구간 합 구하기 5

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

11660번: 구간 합 구하기 5

 

내 코드

import sys
input = sys.stdin.readline()

def make_sum_array(array):
    sum_array = [0] * len(array)
    for i in range(1, len(array)):
        sum_array[i] = sum_array[i - 1] + array[i]
    
    return sum_array


N, M = map(int, input().split())
sum_array = []
for _ in range(N):
    input_array = [0] + list(map(int, input().split()))
    sum_array.append(make_sum_array(input_array))

for _ in range(M):
    result = 0
    x1, y1, x2, y2 = map(int, input().split())
    for x in range(x1 - 1, x2):
        result = result + sum_array[x][y2] - sum_array[x][y1 - 1]
    print(result)

구간 합을 사용했지만 문제에서 요구됐던 시간 제한에 타임 아웃이 걸린 듯하다. 구간 합을 설정하는 부분에서부터 잘못되었다. 나는 행을 기준으로 구간 합을 만들어 완전 탐색 방법과 큰 차이가 없는 것으로 인해 시간 초과가 발생했다고 생각한다.

 

누적 합을 구하는 방법부터 다시 생각해보았다.

 

1

2

3

4

1차원 배열에서 누적합을 구하면 어떻게 될까? 참고로 이 배열은 array라고 하겠다.

 

1

3

6

10

이렇게 나올 것이다. 이걸 좀 다르게 써보면 다음과 같다. 이 배열은 sum_array라고 하겠다. 

 

1

array[1] + sum_array[0]

array[2] + sum_array[1]

array[3] + sum_array[2]

sum_array에 대한 부분을 다시 정의해보면 위와 같다. 1차원 배열 말고 2차원 배열적인 부분으로 넘어가보자.

 

1	2	3	4
2	3	4	5

이것의 누적 합을 구해보면 다음과 같을 것이다.

 

1	3	6	10
3	8	15	24

이 부분에서 나는 누적 합을 잘못 이해하고 풀었던 것이다. 이해하기 쉽게 다시 정리해보겠다.

 

1	1 + 2	1 + 2 + 3	1 + 2 + 3 + 4
1 + 2	1 + 2 + 2 + 3	1 + 2 + 3 + 2 + 3 + 4	1 + 2 + 3 + 4 + 2 + 3 + 4 + 5

그래서 공식이 있을까 가만히 생각해보았다. 사실은 다른 사람이 풀이한 걸 봤는데 이해가 안가서 좀 오래봤다. 8의 경우를 예시르 들어보겠다. 8에 해당하는 인덱스는 (2, 2)이다. array[2][2] + sum_array[1][2] + sum_array[2][1] - sum_array[1][1]을 하게 되면 sum_array[2][2]의 누적 합이 구해진다. 자 이제 코드로 한 번 나타내보자.

 

N, M = map(int, input().split())
array = [[]]
for _ in range(N):
    input_array = [0] + list(map(int, input().split()))
    array.append(input_array)

sum_array = [[0] * (N + 1) for _ in range(N + 1)]

array에 리스트 하나를 더 추가했다. 사실 이 부분은 어떤 값을 넣어도 상관이 없는데 문제에서의 (1, 1) 인덱스를 바로 읽히게 하기 위함으로 추가한 부분이다. (0, 0)은 사용하지 않기 때문에 상관없다.

 

array에는 반복문을 통해 표 데이터를 먼저 받은 다음 append시킨다. 임시로 0을 추가했기 때문에 누적 합을 담을 리스트에도 N + 1 크기의 리스트를 생성해준다.

 

for i in range(1, N + 1):
    for j in range(1, N + 1):
        sum_array[i][j] = array[i][j] # 누적 합의 인덱스에 우선 해당 배열의 값을 초기화한다.
        sum_array[i][j] += sum_array[i - 1][j] + sum_array[i][j - 1]
        sum_array[i][j] -= sum_array[i - 1][j - 1]

누적 합에 대한 리스트를 생성하는 과정이다. 위에서 말했던 것처럼 입력받은 리스트의 i, j번째 인덱스의 값을 먼저 초기화해주고 누적 합의 i - 1, j번째 인덱스와 i, j - 1번째 인덱스 값을 더해주고 i - 1, j - 1번째 인덱스의 값을 빼준다.

 

하지만 이 누적 합에 대한 리스트는 문제에서 원하는 값이 아니다. 3, 4번째 인덱스의 누적 합은 1, 1번째 인덱스부터 3, 4번째 인덱스의 입력 리스트의 값을 더한 것이기 때문이다.

 

for _ in range(M):
    result = 0
    x1, y1, x2, y2 = map(int, input().split())
    result = sum_array[x2][y2]
    result -= sum_array[x1 - 1][y2]
    result -= sum_array[x2][y1 - 1]
    result += sum_array[x1 - 1][y1 - 1]
    print(result)

 

이건 위에서 설명한 추가적인 데이터를 가지고 따지겠다.

 

이렇게 빨간 부분의 구간 합을 구한다고 가정하겠다.

 

그럼 이렇게 각각의 누적 합을 구할 수 있다. 위 표 이미의 구간 인덱스를 보면 (2, 3)부터 (3, 4)까지임을 알 수 있다.(3, 4)의 인덱스 누적 합 값에서 (1, 4)의 누적 합, (3, 2)의 누적 합을 빼고 (1, 2)의 누적 합을 더해보자. 42 - 10 - 15 + 3 = 20으로 구간 합을 구할 수 있었다.

 

여기서 2, 3은 x1, y1이고 3, 4는 x2, y2이다. (1, 4)는 (x1 - 1, y2)이고 (3, 2)는 (x2, y1 - 1)이다. (1, 2)는 (x1 - 1, y1 - 1)이다. 이런 공통점으로 구간합을 구할 수 있다.

 

확실한 건 구간 합에 대한 더 많은 문제를 풀면서 가다듬어봐야겠다.