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

C++로 비트 OR 값이 K 이상인 부분 배열의 개수 구하기

C++를 사용해 비트 OR(bitwise OR) 값이 K 이상인 부분 배열의 개수를 구하는 방법을 알아보겠습니다. 배열 arr[]와 정수 K가 주어졌을 때, 비트 OR 연산 결과가 K보다 크거나 같은 부분 배열(subarray)이 몇 개 존재하는지 찾는 것이 목표입니다.

문제 예시

입력: arr[] = {1, 2, 3}, K = 3
출력: 4

각 부분 배열의 비트 OR:
{1} = 1
{1, 2} = 3
{1, 2, 3} = 3
{2} = 2
{2, 3} = 3
{3} = 3
→ 비트 OR ≥ 3을 만족하는 부분 배열은 총 4개

입력: arr[] = {3, 4, 5}, K = 6
출력: 2

해결 접근 방식

이 문제는 크게 두 가지 방법으로 풀 수 있습니다. 먼저 직관적으로 이해하기 쉬운 완전 탐색 방식을 살펴본 후, 성능을 크게 개선한 효율적인 방식을 다루겠습니다.

1. 브루트 포스 (완전 탐색)

가장 단순한 방법은 만들 수 있는 모든 부분 배열을 하나씩 생성하면서 해당 배열의 OR 값이 K 이상인지 확인하는 것입니다. 조건을 만족하면 정답 카운트를 1씩 증가시킵니다.

예제 코드

#include <bits/stdc++.h>
using namespace std;
int main(){
    int arr[] = {1, 2, 3}; // 주어진 배열.
    int k = 3;
    int size = sizeof(arr) / sizeof(int); // 배열의 크기.
    int answer = 0; // 정답을 세는 카운터 변수.
    for(int i = 0; i < size; i++){
        int bitwise = 0; // k와 비교할 변수.
        for(int j = i; j < size; j++){ // i에서 시작하는 모든 부분 배열.
            bitwise = bitwise | arr[j];
            if(bitwise >= k) // bitwise >= k라면 정답 증가.
                answer++;
        }
    }
    cout << answer << "\n";
    return 0;
}

실행 결과

4

이 방법은 구현이 매우 간단하지만 치명적인 단점이 있습니다. 시간 복잡도가 O(N²)(N은 배열의 크기)이므로, 입력 크기가 커질 경우 실행 시간이 급격히 늘어나 실전 환경에서는 사실상 사용하기 어렵습니다. 이제 더 효율적인 접근 방식을 살펴보겠습니다.

2. 효율적인 접근 — 세그먼트 트리 + 이진 탐색

이 방법은 OR 연산자의 중요한 특성을 활용합니다. OR 연산은 원소를 추가할수록 절대 감소하지 않습니다. 따라서 인덱스 i부터 j까지의 부분 배열의 OR 값이 K 이상이라면, 그 범위를 포함하는 모든 부분 배열의 OR 값 역시 K 이상이 됩니다. 이 성질을 이용하면 각 시작점 i마다 조건을 만족하는 가장 짧은 끝점만 이진 탐색으로 찾으면 됩니다.

구체적으로는 세그먼트 트리로 구간 OR 질의를 빠르게 처리하고, 각 시작 인덱스 i에 대해 이진 탐색으로 OR ≥ K가 처음 되는 끝점 j를 찾습니다. 그러면 j 이후의 모든 끝점에 해당하는 부분 배열도 조건을 만족하므로 한 번에 개수를 더할 수 있습니다.

예제 코드

#include <bits/stdc++.h>
#define N 1000
using namespace std;
int t[4*N];
void build(int* a, int v, int start, int end){ // 세그먼트 트리 생성
    if(start == end){
        t[v] = a[start];
        return;
    }
    int mid = (start + end)/2;
    build(a, 2 * v, start, mid);
    build(a, 2 * v + 1, mid + 1, end);
    t[v] = t[2 * v] | t[2 * v + 1];
}
int query(int v, int tl, int tr, int l, int r){ // 구간 OR 질의 처리
    if (l > r)
        return 0;
    if(tl == l && tr == r)
        return t[v];
    int tm = (tl + tr)/2;
    int q1 = query(2*v, tl, tm, l, min(tm, r));
    int q2 = query((2*v)+1, tm+1, tr, max(tm+1, l), r);
    return q1 | q2;
}
int main(){
    int arr[] = {1, 2, 3}; // 주어진 배열.
    int k = 3;
    int size = sizeof(arr) / sizeof(arr[0]); // 배열의 크기.
    int answer = 0; // 정답을 세는 카운터 변수.
    build(arr, 1, 0, size - 1); // 세그먼트 트리 생성.
    for(int i = 0; i < size; i++){
        int start = i, end = size-1;
        int ind = INT_MAX;
        while(start <= end){ // 이진 탐색
             int mid = (start + end) / 2;
             if(query(1, 0, size-1, i, mid) >= k){ // 부분 배열 조건 검사.
                 ind = min(mid, ind);
                 end = mid - 1;
             }
             else
                 start = mid + 1;
        }
        if(ind != INT_MAX) // 유효한 끝점을 찾았다면 정답 누적.
            answer += size - ind;
    }
    cout << answer << "\n";
    return 0;
}

실행 결과

4

이 접근 방식은 이진 탐색과 세그먼트 트리를 결합해 시간 복잡도를 O(N²)에서 O(N log N)으로 크게 줄여줍니다. 덕분에 앞선 완전 탐색 방식과 달리 입력 크기가 훨씬 큰 경우에도 안정적으로 동작합니다.

마무리

이번 글에서는 이진 탐색과 세그먼트 트리를 활용해 O(n log n) 시간 복잡도로 비트 OR 값이 K 이상인 부분 배열의 개수를 구하는 문제를 해결했습니다. 단순한 완전 탐색 방식과 효율적인 방식 두 가지 접근법을 예제 코드와 함께 살펴봤으며, 동일한 로직은 C, Java, Python 등 다른 언어로도 손쉽게 옮겨 작성할 수 있습니다. 이 글이 문제 풀이에 도움이 되기를 바랍니다.