Post

[Programmers] #181883 - 수열과 구간 쿼리 1 [Java][C++][Python]

정수 배열과 여러 구간이 쿼리로 주어질 때, 각 쿼리 구간의 원소를 모두 1씩 증가시킨 뒤의 배열을 구하는 워밍업 문제.

[Programmers] #181883 - 수열과 구간 쿼리 1 [Java][C++][Python]

문제 링크


1. 아이디어

각 쿼리 [s, e]arr[s]부터 arr[e]까지 모든 원소를 1씩 증가시킨다. 쿼리마다 구간을 직접 순회하며 더해도 배열 길이와 쿼리 수가 모두 최대 $1{,}000$이라 $10^6$ 연산 수준으로 통과한다.

구간이 더 길거나 쿼리가 많은 경우까지 생각하면, 겹치는 구간에서 같은 칸을 반복해 갱신하는 게 낭비다. 이 문제는 모든 쿼리가 구간에 같은 값을 더하는 갱신이고, 갱신이 다 끝난 뒤에야 배열 전체를 한 번 읽는다. 구간 덧셈이 일괄로 들어오고 마지막에 한 번만 조회하는 이 구조가 차분 배열을 그대로 적용할 수 있는 조건이라, 갱신을 쿼리당 $O(1)$로 기록해 두고 마지막에 누적합 한 번으로 복원하면 전체 $O(N + Q)$의 시간복잡도로 해결할 수 있다.


2. 복잡도

접근시간공간
구간 직접 순회$O(N \times Q)$$O(1)$
차분 배열$O(N + Q)$$O(N)$

($N$ = arr의 길이, $Q$ = queries의 길이. 차분 배열 풀이의 diff가 $O(N)$ 공간을 쓴다)


3. 코드

풀이 1: 구간 직접 순회 [Java][C++][Python]

1
2
3
4
5
6
7
8
9
10
11
class Solution {
    public int[] solution(int[] arr, int[][] queries) {
        for (int[] q : queries) {
            for (int i = q[0]; i <= q[1]; i++) {
                arr[i]++;
            }
        }

        return arr;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
#include <bits/stdc++.h>
using namespace std;

vector<int> solution(vector<int> arr, vector<vector<int>> queries) {
    for (auto& q : queries) {
        for (int i = q[0]; i <= q[1]; i++) {
            arr[i]++;
        }
    }

    return arr;
}
1
2
3
4
5
6
def solution(arr, queries):
    for s, e in queries:
        for i in range(s, e + 1):
            arr[i] += 1

    return arr

풀이 2: 차분 배열 [Java][C++][Python]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution {
    public int[] solution(int[] arr, int[][] queries) {
        int n = arr.length;
        int[] diff = new int[n + 1];

        for (int[] q : queries) {
            diff[q[0]]++;
            diff[q[1] + 1]--;
        }

        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += diff[i];
            arr[i] += sum;
        }

        return arr;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#include <bits/stdc++.h>
using namespace std;

vector<int> solution(vector<int> arr, vector<vector<int>> queries) {
    int n = arr.size();
    vector<int> diff(n + 1);

    for (auto& q : queries) {
        diff[q[0]]++;
        diff[q[1] + 1]--;
    }

    int sum = 0;
    for (int i = 0; i < n; i++) {
        sum += diff[i];
        arr[i] += sum;
    }

    return arr;
}
1
2
3
4
5
6
7
8
9
10
from itertools import accumulate


def solution(arr, queries):
    diff = [0] * (len(arr) + 1)
    for s, e in queries:
        diff[s] += 1
        diff[e + 1] -= 1

    return [a + d for a, d in zip(arr, accumulate(diff))]

This post is licensed under CC BY 4.0 by the author.