C++で右側にある自分より小さい要素の数を数える方法

問題の概要

配列 nums が与えられたとき、count[i] に「nums[i] より右側に存在する小さい要素の個数」を格納した配列 count を求める問題を考えてみましょう。

例えば、入力が [5,2,7,1] の場合、結果は [2,1,1,0] となります。

この問題は、Binary Indexed Tree(BIT・フェニック木)を使うことで効率的に解くことができます。単純な二重ループでは O(n²) の計算量が必要になりますが、BIT を利用すれば O(n log m)(m は値の範囲)まで削減できます。

アルゴリズムの手順

• 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 を返します。

• main メソッドでは以下の手順を実行します。

• n := nums のサイズとします。

• サイズ n の配列 res を定義します。

• n が 0 の場合は res をそのまま返します。

• maxx := 0、minn := nums[0] と初期化します。

• i := 1 から i < n の間、i を1ずつ増やしながら次を行います。

  • minn := nums[i] と minn の最小値

• i := 0 から i < n の間、i を1ずつ増やしながら次を行います。

  • nums[i] := nums[i] − minn + 1(すべての値を1以上の正整数に正規化)

  • maxx := nums[i] と maxx の最大値

• サイズ maxx + 1 の配列 bit を定義します。

• i := n − 1 から i >= 0 まで、i を1ずつ減らしながら次を行います。

  • num := nums[i] − 1

  • x := query(num, bit) の呼び出し結果

  • res[i] := x

  • update(num + 1, bit, maxx) を呼び出します。

• res を返します。

実装例

それでは、以下の実装を見て理解を深めましょう。

#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]