n개의 원소로 이루어진 배열 A와 q개의 쿼리 목록 Q가 있다고 가정해 보겠습니다. 각 쿼리 Query[i]는 한 쌍의 값 (x, k)으로 구성되며 다음과 같은 역할을 수행합니다.
- x = 0 (갱신 쿼리): A[k]의 값을 반전합니다. 즉, 1이면 0으로, 0이면 1로 바꿉니다.
- x = 1 (조회 쿼리): 배열에서 k번째로 큰 원소를 출력합니다.
초기 상태에서 배열 A의 모든 원소는 0 또는 1입니다. 이러한 특성 덕분에 매번 배열을 정렬하거나 전체를 탐색할 필요 없이, 배열에 포함된 1의 개수만 추적하면 k번째로 큰 원소가 무엇인지 즉시 판별할 수 있습니다.
예를 들어 입력이 A = [1, 1, 0, 1, 0], Q = [[2, 3], [1, 2], [2, 3], [2, 1], [2, 5]]와 같이 주어졌다면 출력 결과는 [1, 1, 1, 0]이 됩니다.
풀이 접근 방식
이 문제는 다음과 같은 단계를 통해 해결할 수 있습니다.
- 배열 A를 처음부터 끝까지 순회하며 0이 아닌 원소의 개수를 세어 변수 m에 저장합니다. 이 값이 곧 배열 안에 있는 1의 개수입니다.
- 쿼리 목록 Q를 순서대로 순회하며 각 쿼리 (x, k)를 하나씩 처리합니다.
- x가 0이라면 갱신 쿼리입니다. A[k]를 XOR 연산(^= 1)으로 반전한 뒤, 변경 전 값에 따라 m을 조정합니다. 기존 값이 1이었다면 m을 1 감소시키고, 0이었다면 1 증가시킵니다.
- x가 1이라면 조회 쿼리입니다. m이 k보다 크거나 같으면 1을, 그렇지 않으면 0을 출력합니다. 배열이 0과 1로만 구성되어 있으므로 1이 k개 이상 존재한다면 k번째로 큰 원소는 반드시 1이기 때문입니다.
이 알고리즘은 초기화에 O(n), 각 쿼리 처리에 O(1)의 시간이 걸리므로 전체 복잡도는 O(n + q)입니다. 매 쿼리마다 정렬을 수행하는 방식보다 훨씬 효율적입니다.
구현 예제
아래 코드를 통해 실제 구현 과정을 자세히 살펴보겠습니다.
#include <bits/stdc++.h>
using namespace std;
void solve(vector<int> A, vector<vector<int>> Q){
int n = A.size();
int m = 0;
for (int i = 0; i < n; i++){
if (A[i])
m++;
}
for (int j = 0; j < Q.size(); j++){
int x = Q[j][0];
int k = Q[j][1];
if (x == 0){
if (A[k])
m--;
else
m++;
A[k] ^= 1;
}
else{
if (m >= k)
cout << 1 << ", ";
else
cout << 0 << ", ";
}
}
}
int main(){
vector<int> A = { 1, 1, 0, 1, 0 };
vector<vector<int>> Q = { { 1, 2 }, { 0, 1 }, { 1, 2 }, { 1, 0 },{ 1, 4 } };
solve(A, Q);
}
입력
{ 1, 1, 0, 1, 0 }, { { 1, 2 }, { 0, 1 }, { 1, 2 }, { 1, 0 }, { 1, 4 } }
출력
1, 1, 1, 0,
실행 결과 분석
초기 배열에는 1이 세 개 있으므로 m = 3입니다. 첫 번째 쿼리 {1, 2}는 1이 두 개 이상 존재하므로 1을 반환합니다. 두 번째 쿼리 {0, 1}은 인덱스 1의 값을 1에서 0으로 반전하여 m을 2로 만듭니다. 이후 {1, 2}는 1이 정확히 두 개 남아 있으므로 여전히 1을 출력하고, {1, 0} 역시 1을 출력합니다. 마지막 쿼리 {1, 4}는 1이 네 개 존재하지 않으므로 0을 출력합니다. 이처럼 1의 개수만 관리하면 어떤 쿼리든 상수 시간에 처리할 수 있다는 점이 이 풀이법의 핵심입니다.