[백준/C/Python] 11659 - 구간 합 구하기 4

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

구현 과정

C로 구현하면서 의식의 흐름대로 아래처럼 코드를 짰더니 시간 초과가 떴다.

#include <stdio.h>

int main() {
    int n, m;
    scanf("%d %d", &n, &m);

    int num[n];
    for (int i = 0; i < n; i++) {
        scanf("%d", &num[i]);
    }

    while (m--) {
        int i, j;
        scanf("%d %d", &i, &j);

        int res = 0;
        for (int k = i - 1; k < j; k++)
            res += num[k];
        printf("%d\n", res);
    }
    return 0;
}

이 코드에서의 시간 복잡도는 O(N x M)으로 최악의 경우 연산 횟수가 매우 높아진다. 따라서 다른 방법을 사용해야 한다.

누적 합 배열

효율적인 풀이를 위해 누적 합 배열을 사용해야 한다. 이 방법을 사용하면 각 쿼리를 O(1) 시간에 처리할 수 있다. 합을 저장할 배열 sum[n + 1] 을 선언해주고, 합을 미리 계산한 다음에 i, j를 입력받으면서 같이 출력해준다.

'n번째' 수를 쉽게 찾기 위해, 배열의 1번째 요소부터 사용한다. sum[0] = 0 으로 초기화 해준다. i=1부터 n까지 sum[i] = sum[i - 1] + num[i] 로 누적합을 구해 sum 배열에 저장해준다. 이렇게 누적합을 미리 구해놓으면, 구간인 i, j을 입력받을 때 res = sum[j] - sum[i - 1] 로 구간의 합을 빠르게 구할 수 있다.

Python에서의 구현

iter(list) 함수를 사용하면 반복자를 생성할 수 있다. 이게 무슨 말이냐면, 인수로 받은 리스트의 각 요소에 대해 순차적으로 접근할 수 있게 해준다. input = sys.stdin.read 를 써서 모든 입력을 한번에 읽어오는 경우 필요하다.

전체 코드

C

#include <stdio.h>

int main() {
    int n, m;
    scanf("%d %d", &n, &m);

    int num[n + 1];
    int sum[n + 1];
    for (int i = 1; i <= n; i++) {
        scanf("%d", &num[i]);
    }

    sum[0] = 0;
    for (int i = 1; i <= n; i++) {
        sum[i] = sum[i - 1] + num[i];
    }

    while (m--) {
        int i, j;
        scanf("%d %d", &i, &j);

        int res = sum[j] - sum[i - 1];
        printf("%d\n", res);
    }
    return 0;
}

Python

import sys
input = sys.stdin.read

def main():
    data = input().split()
    it = iter(data)

    n = int(next(it))
    m = int(next(it))

    num = [int(next(it)) for _ in range(n)]

    sumList = [0] * (n + 1)

    for i in range(1, n + 1):
        sumList[i] = sumList[i - 1] + num[i - 1]

    res = []
    for _ in range(m):
        i = int(next(it))
        j = int(next(it))

        count = sumList[j] - sumList[i - 1]
        res.append(count)

    print("\n".join(map(str, res)))

if __name__ == "__main__":
    main()