C++でQ番目の人が取得できる棒の最大長を求める方法

この記事では、n本の棒の長さが与えられたとき、順番に棒を受け取っていく各人が取得できる「最大の棒の長さ」を求めるアルゴリズムを、C++のコード例とともに解説します。

問題文

n本の棒の長さが配列として与えられます。ある人が棒を受け取る際には、その時点で最も長い棒の半分((max + 1) / 2)が割り当てられ、残りの部分((max - 1) / 2)は再び利用可能な状態として戻されます。十分な数の棒が常に存在すると仮定し、配列q[]で与えられるM個のクエリに対して、qi番目の人(1から始まる有効な人数)が取得できる最大の棒の長さを答えてください。

例

入力 : a[] = {6, 5, 9, 10, 12}
   q[] = {1, 3}
出力 : 12 9

入力配列が{6, 5, 9, 10, 12}で、クエリ配列が{1, 3}の場合、出力は12と9になります。処理の流れは以下の通りです。

• 1番目の人は、最大の長さである12の棒を取得する
• 配列から12を取り除き、(12 - 1) / 2 = 5 を戻す
• 2番目の人は、次に長い10の棒を取得する
• (10 - 1) / 2 = 4 を戻す
• 3番目の人は、その時点で最大の9の棒を取得する

アルゴリズム

毎回配列全体から最大値を探すのは非効率です。そこで、ソート済みの値を保持するスタックと、切り分けられた残りを保持するキューを組み合わせることで、効率的に解くことができます。

• まずすべての棒の長さをソートし、スタックにプッシュする
• スタックのトップ要素を取り出して結果に記録し、その半分(top / 2)が0でなければキューにプッシュする
• スタックが空の場合は、キューの先頭をポップして結果に記録し、その半分(front / 2)が0でなければキューに戻す
• キューが空の場合は、スタックからポップして結果に記録し、その半分(top / 2)が0でなければキューにプッシュする
• 両方が空でない場合は、スタックのトップとキューの先頭を比較し、大きい方をポップして結果に記録し、その半分をキューに戻す
• スタックとキューの両方が空になるまで繰り返す

こうすることで、各ステップで「現時点で最も長い棒」を常にO(1)で取り出せるようになります。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
vector<int> getMaxRodLength(int *arr, int n, int m) {
    queue<int> q;
    sort(arr, arr + n);
    stack<int> s;
    for (int i = 0; i < n; ++i) {
       s.push(arr[i]);
    }
    vector<int> result;
    while (!s.empty() || !q.empty()) {
       int val;
       if (q.empty()) {
          val = s.top();
          result.push_back(val);
          s.pop();
          val = val / 2;
          if (val) {
            q.push(val);
          }
       } else if (s.empty()) {
          val = q.front();
          result.push_back(val);
          q.pop();
          val = val / 2;
          if (val != 0) {
            q.push(val);
          }
       } else {
          val = s.top();
          int fr = q.front();
          if (fr > val) {
            result.push_back(fr);
            q.pop();
            fr = fr / 2;
            if (fr) {
               q.push(fr);
            }
          } else {
            result.push_back(val);
            s.pop();
            val = val / 2;
            if (val) {
               q.push(val);
            }
          }
       }
    }
    return result;
}
int main() {
    int rods = 5;
    int queries = 10;
    int arr[rods] = {6, 5, 9, 10, 12};
    vector<int> result = getMaxRodLength(arr, rods, queries);
    int query[] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    int n_query = sizeof(query) / sizeof(query[0]);
    cout << "Rod length = ";
    for (int i = 0; i < n_query; ++i) {
       cout << result[query[i] - 1] << " ";
    }
    cout << endl;
    return 0;
}

実行結果

上記のプログラムをコンパイルして実行すると、以下の出力が得られます。

Rod length = 12 10 9 6 6 5 5 4 3 3

この結果から、1番目の人は12、2番目の人は10、3番目の人は9というように、各クエリに対して正しく最大の棒の長さが取得できていることが確認できます。