Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 구간 내 K번째 비트가 설정된 배열 요소 개수 조회하기 — 브루트 포스부터 누적 합 최적화까지

이 글에서는 주어진 범위 내에서 K번째 비트가 설정(set)되어 있는 배열 요소의 개수를 찾는 문제를 다룹니다. 먼저 예시를 살펴보겠습니다.

입력 : arr[] = { 4, 5, 7, 2 }
쿼리 1: L = 2, R = 4, K = 4
쿼리 2: L = 3, R = 5, K = 1
출력 :
    0
    1

이 문제는 먼저 브루트 포스(Brute Force) 방식으로 해결해 보고, 이 방식이 큰 입력 제약 조건에서도 동작하는지 확인해 볼 것입니다. 만약 그렇지 않다면, 더 효율적인 새로운 접근법을 고민하게 될 것입니다.

브루트 포스 접근법

이 방식은 단순히 주어진 범위를 순회하면서 각 요소의 K번째 비트가 설정되어 있는지 확인하고, 설정되어 있다면 카운트를 증가시키는 방법입니다.

예제 코드

#include<bits/stdc++.h>
using namespace std;
#define MAX_BITS 32
bool Kset(int n, int k) { // k번째 비트가 설정되어 있는지 확인
    if (n & (1 << (k - 1)))
    return true;
    return false;
}
int query(int L, int R, int K, int arr[]) {
    int count = 0; // 범위 내 요소 수를 세는 카운터
    for (int i = L; i <= R; i++) { // 범위 순회
        if (Kset(arr[i], K)) {
            count++;
        }
    }
    return count;
}
int main() {
    int arr[] = { 4, 5, 7, 2 }; // 주어진 배열
    int n = sizeof(arr) / sizeof(arr[0]); // 배열의 크기
    int queries[][3] = { // 주어진 L, R, K 값
        { 2, 4, 4 },
        { 3, 5, 1 }
    };
    int q = sizeof(queries) / sizeof(queries[0]); // 쿼리의 개수

    for (int i = 0; i < q; i++) {
        int L = queries[i][0] - 1;
        int R = queries[i][1] - 1;
        int K = queries[i][2];

        cout << query(L, R, K, arr) << "\n";
    }
    return 0;
}

출력 결과

0
1

위 접근법의 시간 복잡도는 O(N*Q)입니다. 여기서 N은 배열의 크기이고 Q는 쿼리의 개수입니다. 보다시피 이 방식은 연산 시간이 너무 많이 걸리기 때문에 높은 제약 조건에는 적합하지 않습니다. 따라서 이제 더 효율적인 접근법으로 프로그램을 작성해 보겠습니다.

효율적인 접근법: 2차원 누적 합(prefix sum) 활용

이 방식에서는 2차원 누적 합(prefix sum) 배열을 유지합니다. 이 배열은 각 인덱스까지 사용된 모든 비트의 개수를 저장하며, 이를 통해 각 쿼리에 대한 답을 O(1) 시간 복잡도로 계산할 수 있습니다.

예제 코드

#include<bits/stdc++.h>
using namespace std;
#define bits 32 // 비트의 개수

int P[100000][bits+1];

bool Kset(int n, int k) {
    if (n & (1 << (k - 1)))
        return true;
    return false;
}
void prefixArray(int n, int arr[]) { // 누적 합 배열 생성
    for (int i = 0; i <= bits; i++) {
        P[0][i] = 0; // 모든 비트의 초기 카운트를 0으로 설정
    }
    for (int i = 0; i < n; i++) {
        for (int j = 1; j <= bits; j++) {
            bool flag = Kset(arr[i], j);
            if (i) // 이전 카운트를 현재 카운트에 더함
                P[i][j] = P[i - 1][j];
            if (flag) { // j번째 비트가 설정되어 있으면 카운트 증가
                P[i][j]++;
            }
        }
    }
}
int query(int L, int R, int K) {
    if (L) // L이 0이 아니라면 R까지의 누적 합에서 L-1까지의 누적 합을 뺌
        return P[R][K] - P[L - 1][K];
    else
        return P[R][K];
}
int main() {
    int arr[] = { 8, 9, 1, 3 }; // 주어진 배열
    int n = sizeof(arr) / sizeof(arr[0]); // 배열의 크기
    int queries[][3] = {
        { 1, 3, 4 },
        { 2, 4, 1 }
    };
    prefixArray(n, arr); // 누적 합 배열 생성 함수 호출
    int q = sizeof(queries) / sizeof(queries[0]); // 쿼리의 개수

    for (int i = 0; i < q; i++) {
        int L = queries[i][0] - 1;
        int R = queries[i][1] - 1;
        int K = queries[i][2];
        cout << query(L, R, K) << "\n";
    }
    return 0;
}

출력 결과

2
3

누적 합 배열을 활용하면 O(1)만에 답을 구할 수 있으므로, 전체 시간 복잡도는 O(N)으로 대폭 감소합니다. 여기서 N은 주어진 배열의 크기입니다.

코드 설명

이 프로그램에서는 배열의 모든 인덱스에 대해 해당 인덱스까지 각 비트가 몇 번 등장했는지를 세는 누적 카운터(prefix counter)를 유지합니다. 배열에 대한 누적 카운트를 미리 구축해 두면, K번째 비트의 구간 내 개수를 구할 때 R 인덱스까지의 K번째 비트 누적 개수에서 L-1 인덱스까지의 누적 개수를 빼면 곧바로 정답을 얻을 수 있습니다.

예를 들어 배열 {8, 9, 1, 3}의 경우:

  • 쿼리 (1, 3, 4): 인덱스 0~2 범위에서 4번째 비트가 설정된 요소는 8(1000₂)과 9(1001₂)로 2개입니다.
  • 쿼리 (2, 4, 1): 인덱스 1~3 범위에서 1번째 비트가 설정된 요소는 9(1001₂), 1(0001₂), 3(0011₂)로 3개입니다.

결론

이 글에서는 구간 내 K번째 비트가 설정된 배열 요소의 개수를 구하는 쿼리 문제를 해결했습니다. 일반적인 브루트 포스 방식과 누적 합을 활용한 효율적인 방식 두 가지 접근법과 함께 C++ 프로그램도 살펴보았습니다. 동일한 로직은 C, Java, Python 등 다른 프로그래밍 언어로도 손쉽게 구현할 수 있습니다. 특히 쿼리가 많은 상황에서는 누적 합 기반 접근법이 압도적으로 유리하다는 점을 기억해 두시길 바랍니다. 이 글이 도움이 되었기를 바랍니다.