C++で2つ目の配列を並べ替えて作れる辞書式順序最小のシーケンスを求める方法

2つの配列 A と B(いずれも n 個の整数を持つ)が与えられたとします。ここで、B の要素をうまく並べ替えることで、(A[i] + B[i]) % n で構成されるシーケンスが辞書式順序で最小になるようにします。最終的に、その最小のシーケンスを返すのが目的です。

たとえば、入力が A = {1, 2, 3, 2}、B = {4, 3, 2, 2} の場合、出力は [0, 0, 1, 2] となります。

解法の考え方

この問題は貪欲法(greedy法)で解くことができます。各位置 i において、(a[i] + b) % n が最小になるような b を残っている候補の中から選んでいく方針です。

(a[i] + b) % n が最小値 0 になるのは、b = (n − a[i]) のときです。もしその値が存在しない場合は、「n − a[i]」以上の最小の要素を選び、それもなければセット内の最小の要素に折り返して選択します。この探索を効率化するために、std::set の lower_bound() を利用すると、各要素を O(log n) で見つけられます。

アルゴリズムの手順

• n := 配列 a のサイズとする
• マップ my_map(各値の出現回数を管理)を定義する
• セット my_set(使用可能な値の集合)を定義する
• i が 0 から n 未満の間、以下を繰り返す:
  • my_map[b[i]] を 1 増やす
  • b[i] を my_set に挿入する
• 結果を格納する配列 sequence を定義する
• i が 0 から n 未満の間、以下を繰り返す:
  • a[i] が 0 の場合:
    • it := my_set 内で 0 以上の最初の要素(lower_bound(0))
    • value := *it
    • value % n を sequence の末尾に追加する
    • my_map[value] を 1 減らし、0 になったら my_set から value を削除する
  • それ以外の場合:
    • x := n − a[i]
    • it := my_set 内で x 以上の最初の要素
    • 該当する要素が存在しない場合(it == end)、it := my_set の先頭(最小要素)に戻す
    • value := *it
    • (a[i] + value) % n を sequence の末尾に追加する
    • my_map[value] を 1 減らし、0 になったら my_set から value を削除する
• sequence を返す

実装例(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;
}
vector<int> solve(vector<int>&a, vector<int> &b) {
   int n = a.size();
   unordered_map<int, int> my_map;
   set<int> my_set;
   for (int i = 0; i < n; i++) {
      my_map[b[i]]++;
      my_set.insert(b[i]);
   }
   vector<int> sequence;
   for (int i = 0; i < n; i++) {
      if (a[i] == 0) {
         auto it = my_set.lower_bound(0);
         int value = *it;
         sequence.push_back(value % n);
         my_map[value]--;
         if (!my_map[value])
            my_set.erase(value);
      }
      else {
         int x = n - a[i];
         auto it = my_set.lower_bound(x);
         if (it == my_set.end())
            it = my_set.lower_bound(0);
         int value = *it;
         sequence.push_back((a[i] + value) % n);
         my_map[value]--;
         if (!my_map[value])
            my_set.erase(value);
      }
   }
   return sequence;
}
int main() {
   vector<int> a = {1, 2, 3, 2};
   vector<int> b = {4, 3, 2, 2};
   vector<int> res = solve(a, b);
   print_vector(res);
}

入力

{1, 2, 3, 2}, {4, 3, 2, 2}

出力

[0, 0, 1, 2]

計算量について

各要素の挿入・削除・探索は std::set 上で O(log n) に行われるため、全体の時間計算量は O(n log n) となります。また、マップとセットにはそれぞれ O(n) の追加メモリが必要です。貪欲に各位置で可能な限り小さい剰余を選ぶことで、結果として得られるシーケンス全体が辞書式順序で最小になることが保証されます。