F - Transportation Expenses Editorial by MMNMM

線形時間の解法

まず、次の \(\Theta(N\log N)\) 時間の解法について説明します。

  1. \(A\) をソートし、\(A _ 1\leq A _ 2\leq\dotsb\leq A _ N\) とする。
  2. はじめ \(S=0\) とし、\(i=1,2,\ldots,N\) の順に次の手順を行う。
    1. \(S+A _ i\times(N-i+1)\) が \(M\) より大きいなら、繰り返しを終了する。
    2. \(S\leftarrow S+A _ i\) とする。
  3. 繰り返しが最後まで続いた場合、infinite と出力し、終了する。
  4. そうでない場合、繰り返しを終了した時点での \(i\) について \(\left\lfloor\dfrac{M-S}{N-i+1}\right\rfloor\) を出力し、終了する。

これは、交通費を満額支給される人数を全探索していると考えることができます。
\(A _ 1,A _ 2,\dotsc,A _ N\) について、その金額を上限としたときに予算内に収まるかどうかを判定し、\(A _ {i-1}\) なら収まり \(A _ i\) では収まらないような \(i\) を見つけることで満額支給される人数を決定することができます(便宜上 \(A _ 0=0\) とします)。

この解法はソートがボトルネックとなり \(\Theta(N\log N)\) 時間となっています。

しかし、答えを求めるために列 \(A\) は必ずしもソートされている必要はありません。 列 \(A\) を並べ替えることで、ある \(i\ (1\leq i\leq N)\) について

  • \(A _ j\leq A _ {i-1}\ (j\lt i)\)
  • \(A _ {i-1}\leq A _ i\)
  • \(A _ i\leq A _ j\ (i\leq j)\)
  • \(\displaystyle\sum _ {j\lt i-1}A _ j+(N-i+2)A _ {i-1}\leq M\lt\sum _ {j\lt i}A _ j+(N-i+1)A _ i\)

が成り立つようにできれば、交通費を満額支給される人数が \(i-1\) 人であることがわかります。

このような条件を満たす列を \(O(N)\) 時間で構成する方法について考えます。


Quick Select と呼ばれるアルゴリズムがあります。 これは、ソートされていない数列 \(A\) に対して、小さいほうから \(k\) 番目の値を選び、列を \(A _ j\leq A _ k\ (j\leq k)\) かつ \(A _ k\leq A _ j\ (k\leq j)\) を満たすように並べ替えるアルゴリズムです。

Quick Select は以下のようなアルゴリズムです。

  1. 列の左端と右端を表す変数 \(L,R\) を \(L=0,R=N\) と、選ぶ対象を表す値 \(x\) を \(x=k\) と初期化する。
  2. \(L\lt R\) である限り、以下を繰り返す。
    1. \(A _ L,A _ {L+1},\dotsc,A _ {R-1}\) から要素をひとつ選び、その値を \(p\) とする。
    2. \(A _ L,A _ {L+1},\dotsc,A _ {R-1}\) を \(p\) 未満の値と \(p\) と等しい値と \(p\) より大きい値の \(3\) つに分け、この順に連結する。
    3. それぞれの個数を \(C _ {\operatorname{less}},C _ {\operatorname{equal}},C _ {\operatorname{greater}}\) とし、次の条件分岐を行う。
      1. \(x\leq C _ {\operatorname{less}}\) ならば、\(R\leftarrow L+C _ {\operatorname{less}}\) とする。
      2. \(C _ {\operatorname{less}}\lt x\leq C _ {\operatorname{less}}+C _ {\operatorname{equal}}\) ならば、繰り返しを終了する。
      3. \(C _ {\operatorname{less}}+C _ {\operatorname{equal}}\lt x\) ならば、\(L\leftarrow R-C _ {\operatorname{greater}},x\leftarrow x-C _ {\operatorname{less}}-C _ {\operatorname{equal}}\) とする。

このアルゴリズムは、2-1 での \(p\) の選び方によって計算量が average \(O(N)\) や expected \(O(N)\) や worst \(O(N)\) に変化します。 適切に選ぶと worst \(O(N)\) とすることができます(少し実装が煩雑になり、計算量の定数倍も大きくなります)が、事前に \(A\) をランダムにシャッフルしておくと、\(A _ L\) を \(p\) として選ぶことで expected \(O(N)\) にすることができます。

Quick Select をもとに 2-3 の場合分けの部分を適切に変更することで、今回の問題を解くことができます。
2-3 で適切に判定を行うために、\(A _ 0,\ldots,A _ {L-1}\) の総和などを追加で管理する必要があります。

実装例は以下のようになります。 以下の実装例では事前に \(A\) をシャッフルすることで時間計算量を expected \(O(N)\) としていますが、適切にピボットを選択することで worst \(O(N)\) とすることもできます。

#include <iostream>
#include <vector>
#include <algorithm>
#include <random>

int main() {
    using namespace std;
    unsigned N;
    unsigned long M;
    cin >> N >> M;
    vector<unsigned> A(N);
    for(auto&& a : A)
        cin >> a;

    // A をシャッフル (平均計算量を期待計算量に)
    {
        mt19937_64 rng{random_device{}()};
        ranges::shuffle(A, rng);
    }

    // 列の左端・右端
    auto first = begin(A), last = end(A);
    // prefix_sum := 先頭から first までの総和
    // suffix_count := last から末尾までの個数
    unsigned long prefix_sum = 0, suffix_count = 0;
    while(first != last){
        // pivot を選び、区間 [first, last) を [より小さい][等しい][より大きい] の順に並べる
        // それぞれの境界を mid_first, mid_last とする
        const auto pivot = *first;
        auto mid_first = partition(first, last, [pivot](const auto a){return a < pivot;});
        auto mid_last = partition(mid_first, last, [pivot](const auto a){return a <= pivot;});

        // 上限を pivot にしたときの合計金額を求める
        const auto tmp_prefix_sum = prefix_sum + reduce(first, mid_first, 0UL);
        const auto tmp_suffix_count = suffix_count + distance(mid_last, last);
        const auto pivot_count = distance(mid_first, mid_last);
        const auto total_subsidy = tmp_prefix_sum + pivot * (pivot_count + tmp_suffix_count);

        // 予算に収まっているかどうかを判定
        if (total_subsidy <= M) {
            first = mid_last;
            prefix_sum = tmp_prefix_sum + pivot * pivot_count;
        } else {
            last = mid_first;
            suffix_count = pivot_count + tmp_suffix_count;
        }
    }

    // 全員に払えるなら infinite
    if (last == end(A)) {
        cout << "infinite" << endl;
        return 0;
    }

    // 上限を超えた人たちに均等に配ったものが答え
    cout << (M - prefix_sum) / suffix_count << endl;
    return 0;
}

posted:
last update: