情報系院生の活動日誌

いまはもう社会人

ABC149 E - Handshake:「~は累積和を用いることで O(1) で求まるので」ではわからない人のための累積和解説

atcoderの解説は往々にしてストロングスタイルなところが好きです。

atcoder.jp   本番さっぱりとけなかったので、後日解説を読みました。

まず、増える幸福度が、ある定数 X 以上になる握手の方法が何通りあるか、を求めることを考えます。

左手で握手する人を先に決めたとき、右手で握手できる人が何人いるかは累積和を用いることでO(1)で求まるのでトータルで O(N)で計算できます。

ここが全然わからなくて、解説の簡潔さを恨みながらようやく理解したので備忘録に残しておきます。


まず、Nの制約が105であるので、一つ一つの握手の幸福度を考えていると間に合わないことがわかります。

ここで、解説にある通り発想を変えて、増える幸福度が定数X以上になる握手の方法が何通りあるかを考えます。しかし、一つ一つの握手の幸福度をそれぞれ見ていくと間に合わないので、「累積和」と呼ばれる手法を使います。

qiita.com

この記事は簡単な累積和の解説記事ですが、僕はこの問題にどのように累積和を適用すればいいのかわからなくてそのまま年を越しました。

ところで、「いもす法」というアルゴリズムがあります。いもす法とは、累積和のアルゴリズムを多次元、多次数に拡張したものです。

imoz.jp

こちらの原典の最初の例を読むとなんとなくこの問題への適用方法がわかってきました。例えば、

  1. 一回の握手の幸福度の取りうる値域(<= 2 * 105)分のサイズの配列を用意。
  2. 配列を0で初期化。
  3. それぞれの握手の幸福度をvとしたとき、array[v]を+1。
  4. その配列に対して後ろから累積和を計算。

とすることで、ある値Xに対して、「array[X] = X以上になる握手の方法の数」となるような配列arrayが得られます(考え方はリンク先の現在入店している人数を数える場合と同じ)。

この配列を用いれば、「array[X] = X以上になる握手の方法の数」がO(1)で求まります。

しかしこのような配列を作成するのに全通りの握手の幸福度を調べているので、配列の構築に計算量がO(N2)かかってしまいます。

そこで、上記のアルゴリズムを片手の分だけで行い、累積和を計算します。

このようにすると、例えば左手が握るゲストの手を固定し、左手側のゲストの値がkのとき、右手側の値がX-k以上になるゲストの数を累積和を用いることで O(1) で求めることができます。 したがって、固定したゲストを全通り試すことで、トータルでO(n)でX以上になる握手の方法の数が求まります。

あとは、このXの値を二分探索で求めます。X以上になる握手の方法の数、がちょうどMになるとき、そのM通りの握手の幸福度を合計したものが答えです。

X以上になる握手の方法の数がちょうどMにならない場合もあるので、その場合は多少考えることが増えますが、その部分は以下の記事がとてもわかりやすかったです。

yamakasa.net