[알고리즘] 구간 합 (Prefix Sum)

나영·2023년 8월 14일

알고리즘

목록 보기
5/8
post-thumbnail

구간 합 vs 부분 합

  • 구간 합 (Prefix sum) : 특정 구간의 합 (보통 1차원 배열에서 i ~ k 인덱스 사이의 값들의 합)
  • 부분 합 (Partial sum) : 처음부터 특정 인덱스까지의 합 (0 ~ k 인덱스 사이의 값들의 합)

구간 합 알고리즘

  • 반복문을 사용해 i ~ k 까지의 값을 더하게 되면 시간복잡도는 O(n) 이다. 하지만 n 의 값이 크면 시간초과가 기 마련이다.
  • 구간 합 알고리즘을 사용하면 O(1) 의 성능을 갖는다.
  • sum 배열을 활용해 sum[k] - sum[i - 1] 로 표현한다.

예제 1

백준 구간 합 구하기 4 문제

class Main_11659 {
    static int[] arr;
    static int ans;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken()); // 수의 개수
        int m = Integer.parseInt(st.nextToken()); // 합 구해야하는 횟수

        arr = new int[n + 1]; // 누적합 저장할 배열
        arr[0] = 0; // 0번째 인덱스는 0으로 초기화
		
        // 누적합 배열에 저장 
        st = new StringTokenizer(br.readLine());
        for(int i = 0; i < n; i++){
            arr[i + 1] = arr[i] + Integer.parseInt(st.nextToken());
        }

		// 구간합 구하기 
        for(int i = 0; i < m; i++){
            st = new StringTokenizer(br.readLine());
            int first_idx = Integer.parseInt(st.nextToken());
            int last_idx = Integer.parseInt(st.nextToken());

            accumulate(first_idx, last_idx);
            System.out.println(ans);
        }
    }

    // 구간합 구하는 함수 
    public static void accumulate(int first_idx, int last_idx){
        ans = arr[last_idx] - arr[first_idx - 1];
    }
}

O(NM) -> O(N + M) 으로 시간복잡도를 줄일 수 있다.

예제 2

백준 구간 합 구하기 5 문제

class Main_11660 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        int n = Integer.parseInt(st.nextToken()); // 행렬 크기
        int m = Integer.parseInt(st.nextToken()); // 답 구해야 할 횟수
        int[][] arr = new int[n + 1][n + 1]; // 주어진 수 배열
        int[][] dp = new int[n + 1][n + 1]; // 누적합 배열

        // arr 배열에 주어진 값 넣기
        for(int i = 1; i <= n; i++){
            st = new StringTokenizer(br.readLine());

            for(int j = 1; j <= n; j++){
                arr[i][j] = Integer.parseInt(st.nextToken());
            }
        }

        // dp 배열에 누적합 넣기
        for(int i = 1; i <= n; i++){
            for(int j = 1; j <= n; j++){
                dp[i][j] = dp[i - 1][j] + dp[i][j - 1] - dp[i -1][j - 1] + arr[i][j];
            }
        }

        // x1, y1, x2, y2
        for(int i = 0; i < m; i++){
            st = new StringTokenizer(br.readLine());
            int x1 = Integer.parseInt(st.nextToken());
            int y1 = Integer.parseInt(st.nextToken());
            int x2 = Integer.parseInt(st.nextToken());
            int y2 = Integer.parseInt(st.nextToken());

            int ans = dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1];
            sb.append(ans + "\n");
        }

        System.out.print(sb);

    }
}
  • 누적합 배열에 넣기
    dp[i][j] = dp[i - 1][j] + dp[i][j - 1] - dp[i -1][j - 1] + arr[i][j];

  • 구간합 구하기
    dp[x2][y2] - dp[x1 - 1][y2] - dp[x2][y1 - 1] + dp[x1 - 1][y1 - 1];

그림 출처


⭐ 정리

앞서, 투포인터와 슬라이딩 윈도우에 관해서도 정리한 글이 있다.
3가지 알고리즘의 사용 방법을 정리해보자면,

  • 구간의 길이가 가변적인 경우 -> 투포인터
  • 구간의 길이가 고정된 경우 -> 슬라이딩 윈도우
  • 쿼리를 통해 구간 합여러 번 구해야하는 경우 -> 구간 합

2개의 댓글

comment-user-thumbnail
2023년 8월 14일

글 재미있게 봤습니다.

1개의 답글