C++でターゲットの色までの最短距離を求めるアルゴリズムを解説

問題の概要

1、2、3の3種類の色が格納された配列colorsがあるとします。いくつかのクエリが与えられ、各クエリは2つの整数iとcから構成されます。このとき、指定されたインデックスiからターゲットの色cまでの最短距離を求める必要があります。該当する色が存在しない場合は-1を返します。

例えば、colors配列が[1,1,2,1,3,2,2,3,3]、queries配列が[[1,3],[2,2],[6,1]]である場合、出力は[3,0,3]となります。その理由は以下の通りです。

• インデックス1から最も近い「3」はインデックス4に存在するため、距離は3
• インデックス2から最も近い「2」はインデックス2自身であるため、距離は0
• インデックス6から最も近い「1」はインデックス3に存在するため、距離は3

解法のアプローチ

この問題を効率的に解くには、あらかじめ各色の出現位置を記録しておき、クエリごとに二分探索(lower_bound)を使って最も近い位置を特定します。手順は以下の通りです。

• 4行の行列indexを作成します(nはcolors配列の要素数)
• iを0からn-1までループし、index[colors[i]]にiを挿入して各色のインデックスを記録します
• 各クエリに対して以下を処理します
  • x := queries[i][0]、c := queries[i][1]とします
  • index[c]が空の場合、retに-1を挿入して次のクエリへスキップします
  • lower_boundを使い、index[c]の中からx以上となる最初の要素の位置itを求めます
  • op1 := 無限大、op2 := 無限大で初期化します
  • itがindex[c]のサイズと等しい場合(x未満の要素しかない場合)、itを1減らしてop1 := |x - index[c][it]|を計算します
  • itが0の場合は、op1 := |x - index[c][it]|を計算します
  • それ以外の場合は、op1 := |x - index[c][it]|を計算し、itを1減らしてからop2 := |x - index[c][it]|も計算します
  • op1とop2の最小値をretに挿入します
• 最後にretを返します

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:
   vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) {
      vector < vector <int> >idx(4);
      int n = colors.size();
      for(int i = 0; i < n; i++){
         idx[colors[i]].push_back(i);
      }
      vector <int> ret;
      for(int i = 0; i < queries.size(); i++){
         int x = queries[i][0];
         int c = queries[i][1];
         if(idx[c].size() == 0){
            ret.push_back(-1);
            continue;
         }
         int it = lower_bound(idx[c].begin(), idx[c].end() , x) - idx[c].begin();
         int op1 = INT_MAX;
         int op2 = INT_MAX;
         if(it == idx[c].size()){
            it--;
            op1 = abs(x - idx[c][it]);
         }
         else if(it == 0){
            op1 = abs(x - idx[c][it]);
         }
         else{
            op1 = abs(x - idx[c][it]);
            it--;
            op2 = abs(x - idx[c][it]);
         }
         ret.push_back(min(op1, op2));
      }
      return ret;
   }
};
main(){
   vector<int> v = {1,1,2,1,3,2,2,3,3};
   vector<vector<int>> v1 = {{1,3},{2,2},{6,1}};
   Solution ob;
   print_vector(ob.shortestDistanceColor(v, v1));
}

入力

[1,1,2,1,3,2,2,3,3]
[[1,3],[2,2],[6,1]]

出力

[3,0,3]

計算量について

前処理(各色のインデックス記録)にはO(n)の時間がかかり、各クエリの処理は二分探索によりO(log n)で完了します。したがって、クエリ数をqとすると、全体の計算量はO(n + q log n)となり、クエリが多い場合でも効率的に動作します。各色の位置を事前にソート済みのリストとして保持しておくことが、このアルゴリズムの鍵となります。