AtCoderログ：0009 - ABC 208 C

問題

問題文

高橋王国には $N$ 人の国民がいます。 全ての国民には国民番号が割り振られており、$i$ 人目の国民の国民番号は $a_i$ です。ここで、$a_i$ は互いに異なります。
高橋君は $K$ 個のお菓子を持っています。高橋君は次のルールに従って、持っているお菓子が無くなるまで国民にお菓子を配ることにしました。
・持っているお菓子が $N$ 個以上ある場合、全員に $1$ 個ずつお菓子を配る。
・そうでない場合、その時点で高橋くんが持っているお菓子の個数を $K′$ として、国民番号が小さい方から $K′$ 人に $1$ 個ずつ配る。より厳密には、$a_i$ の値が小さい方から $K′$ 人を選び、選んだ人に $1$ 個ずつお菓子を配る。
全てのお菓子を配り終えたとき、$i$ 人目の国民は何個のお菓子を持っていますか？

制約

・$1 \le N \le 2 \times 10^5$
・$1 \le K \le 10^{18}$
・$1 \le a_i \le 10^9$
・$a_i$ は互いに異なる。
・入力は全て整数である。

回答

回答1 (AC)

国民の人数を n 人、お菓子の個数を k 個とするとき、処理の流れは以下のようになります。
(1) 持っているお菓子が n 個以上の場合：国民全員に同じ個数ずつのお菓子を配る。配るお菓子の個数は k/n 個ずつで、配った後のお菓子の個数は k%n (=k-(k/n)*n) 個になる。
(2) そうでない場合：国民番号をソートして小さい方から $K'$(=k%n) 番目までの国民を特定し、その国民にお菓子を 1 個ずつ配る。
(1)は簡単にコーディングできます。(2)では国民番号 $a_i$ からインデックス $i$ を求める操作が必要です。

国民番号を保持する配列を a(n), ソートするために保持する国民番号の配列を a_sorted(n), 配布されるお菓子の個数を保持する配列を sweet(n) とします。
(2)を実現するのに最も簡単なのは、さらに逆引き用の配列 id_to_index(k) を使用し、 id_to_index.at(a_i)=i と設定することですが、今回は $1 \le K \le 10^{18}$ であるため、逆引き用の配列のメモリを確保するのが厳しそうです。
そこで国民番号からインデックスを効率的に逆引きするために、連想配列 map (Python を辞書型) 使ってみました。仮想配列 index_to_id[] により逆引きを実現しています。コードは以下のようになりました。なお、私が ABC 208 に参加したときの提出コードはこの方針で作成しました。

abc208c-1.cpp

#include <bits/stdc++.h>
using namespace std;
 
int main() {
  int64_t n, k;
  cin >> n >> k;
 
  vector<int64_t> a(n,0), a_sorted(n,0);
  map<int64_t,int> id_to_index;
  for ( int i=0; i<n; i++ ) { 
    cin >> a.at(i);
    id_to_index[a.at(i)] = i;
  }
  copy( a.begin(), a.end(), a_sorted.begin() );

  sort( a_sorted.begin(), a_sorted.end() );
  vector<int64_t> sweet(n,k/n);
  k = k%n;
  for ( int i=0; i<k; i++ ) {
    sweet.at(id_to_index[a_sorted.at(i)]) += 1;
  }
 
  for ( int i=0; i<n; i++ ) {
    cout << sweet.at(i) << endl;
  }
}

回答2 (AC)

回答1の出力を見て気づくことは、出力されるお菓子の個数は(高々)2種類しかないことです。従って、(1)で計算した全員に配るお菓子の個数だけを保持していれば、配列 sweet(n) は必要ないことがわかります。また今回の問題では、追加のお菓子が配られるかどうかは小さい順にソートした際の $K'$ 番目の国民番号が境になっているので、$K'$ 番目の国民番号 id さえわかっていれば、追加のお菓子を配るかどうかの判定が可能となります。つまり仮想配列 id_to_index[] を使用する必要がありません。これらの考察を適用したコードは以下のようになりました。

abc208c-2.cpp

#include <bits/stdc++.h>
using namespace std;
 
int main() {
  int64_t n, k;
  cin >> n >> k;
 
  vector<int64_t> a(n,0), a_sorted(n,0);
  for ( int i=0; i<n; i++ ) { 
    cin >> a.at(i);
  }
  copy( a.begin(), a.end(), a_sorted.begin() );
 
  int64_t sweet = k/n;
  k = k%n;
  sort( a_sorted.begin(), a_sorted.end() );
  int64_t id = a_sorted.at(k);
  for ( int i=0; i<n; i++ ) {
    if ( a.at(i)<id ) {
      cout << sweet+1 << endl;
    } else {
      cout << sweet << endl;
    }
  }
}

調べたこと

AtCoder の解説 → 公式解説

{a_i} をソートするのではなく、{i} をソートする方法が紹介されています。ペア (i,a(i)) を2項目によってソートするイメージですかね。勉強になります。

AtCoder の解説 → ユーザ解説

はやりインデックスの方をソートしています。常識なのでしょう。あと、(1)の処理が難しい旨が指摘されていますが、そうなんでしょうか (私は全く引っかからなかったので)。

学んだこと

• 仮想配列 map の使い方 (宣言、代入)
• 配列の使い方 (コピー)
• ${a_1,a_2,\dots,a_n}$ の値をもとに ${1, 2,\dots,n}$ をソートする方法

リンク

ABC 208 関連情報

• AtCoderログ：0006 - ABC 208 に参加しました
• AtCoderログ：0007 - ABC 208 A
• AtCoderログ：0008 - ABC 208 B
• AtCoderログ：0009 - ABC 208 C

前後の記事

• 前の記事 → AtCoderログ：0008 - ABC 208 B
• 次の記事 → [AtCoderログ：0010 - ABC 081 A]
  (https://qiita.com/RoadLynton/items/dafadf919a845989d44d)