C++로 풀어보는 '자기 자신 뒤의 더 작은 숫자 개수' 문제

문제 소개

배열 nums가 주어졌을 때, 각 원소 기준으로 자신의 오른쪽에 있는 더 작은 숫자의 개수를 저장하는 새로운 배열 count를 구하는 것이 목표입니다. 즉, count[i]에는 nums[i]보다 작으면서 인덱스가 i보다 뒤에 있는 원소들의 개수가 들어갑니다.

예를 들어 입력이 [5,2,7,1]이라면 결과는 [2,1,1,0]이 됩니다.

5의 오른쪽에는 2와 1, 두 개의 더 작은 숫자가 있습니다.
2의 오른쪽에는 1 하나만 있습니다.
7의 오른쪽에는 1 하나가 있습니다.
1의 오른쪽에는 아무 숫자도 없으므로 0입니다.

접근 방법: 펜윅 트리(Fenwick Tree)

이 문제는 단순한 이중 반복문으로 풀면 O(n²)의 시간이 걸리지만, 바이너리 인덱스드 트리(BIT, 펜윅 트리)를 활용하면 O(n log n)으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 배열을 오른쪽에서 왼쪽으로 순회하면서, 지금까지 등장한 값들의 빈도를 BIT에 기록하고 현재 값보다 작은 값의 개수를 질의(query)하는 것입니다.

알고리즘 단계

update() 메서드 정의: 인덱스, bit 배열, n을 매개변수로 받습니다.
index <= n인 동안 다음을 반복합니다.
- bit[index] 값을 1 증가시킵니다.
- index = index + (index AND -index)로 갱신합니다.
query() 메서드 정의: 인덱스와 bit 배열을 매개변수로 받습니다.
- ans := 0으로 초기화합니다.
- index > 0인 동안 다음을 반복합니다.
  - ans = ans + bit[index]
  - index = index - (index AND -index)
- ans를 반환합니다.
메인 로직:
n := nums의 크기로 설정하고, 크기 n의 결과 배열 res를 선언합니다.
n이 0이면 그대로 res를 반환합니다.
maxx := 0, minn := nums[0]으로 초기화한 뒤, 배열을 순회하며 최솟값 minn을 구합니다.
모든 원소를 nums[i] = nums[i] - minn + 1로 변환하여 값을 양수 범위로 정규화(normalization)하고, 최댓값 maxx를 갱신합니다. 이 과정은 음수가 포함된 입력도 BIT의 인덱스로 사용할 수 있게 해줍니다.
크기가 maxx + 1인 bit 배열을 생성합니다.
i = n-1부터 0까지 역순으로 순회하면서:
- num = nums[i] - 1로 설정합니다.
- x := query(num, bit)를 호출하여 현재 값보다 작은 원소의 개수를 구합니다.
- res[i] := x로 저장합니다.
- update(num + 1, bit, maxx)를 호출하여 현재 값을 BIT에 기록합니다.
최종적으로 res를 반환합니다.

C++ 구현 예제

아래 코드를 통해 실제 구현을 확인해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
class Solution {
   public:
   void update(int index,vector <int>& bit, int n){
      while(index<=n){
         bit[index]++;
         index += (index)&(-index);
      }
   }
   int query(int index, vector <int> bit){
      int ans = 0;
      while(index>0){
         ans+=bit[index];
         index -= index&(-index);
      }
      return ans;
   }
   vector<int> countSmaller(vector<int>& nums) {
      int n = nums.size();
      vector <int> res(n);
      if(!n)return res;
      int maxx = 0;
      int minn = nums[0];
      for(int i =1;i<n;i++)minn = min(nums[i],minn);
      for(int i =0;i<n;i++){
         nums[i] = nums[i]-minn+1;
         maxx = max(nums[i],maxx);
      }
      vector <int> bit(maxx+1);
      for(int i =n-1;i>=0;i--){
         int num = nums[i]-1;
         int x = query(num,bit);
         res[i] = x;
         update(num+1,bit,maxx);
      }
      return res;
   }
};
main(){
   Solution ob;
   vector<int> v = {5,2,7,1};
   print_vector(ob.countSmaller(v));
}

입력

[5,2,7,1]

출력

[2, 1, 1, 0]

마무리

펜윅 트리를 사용하면 각 원소마다 오른쪽에 있는 더 작은 숫자의 개수를 로그 시간 안에 계산할 수 있어, 전체 시간 복잡도가 O(n log n)으로 줄어듭니다. 값의 범위를 정규화하는 전처리 단계 덕분에 음수나 큰 수가 섞여 있는 입력에도 안정적으로 동작한다는 점도 기억해 두면 좋습니다.