は?「計算結果がK番目に大きくなる組を求める」? 何を言っているの?と思うかもしれませんが、こういうことです。

D – Cake 123

D – Cake 123

問題の趣旨は「配列 A, B, C が与えられる。A[i] + B[j] + C[k] が大きい順に K 個出力せよ」というものです。

ただし配列の長さは最大 1,000 なので全パターン(10億通り)試して上位 K 個を取るという解法は使えません。

こんなときに使えるのが優先度付きキュー(priority queue)です。優先度付きキューは格納するデータに優先度を設定して優先度の値が小さい順に取り出すことができるデータ構造です。取り出すことができるのは先頭の値だけですが、全体をソートするときの時間計算量が O(NlogN) であることに対して優先度付きキューであれば O(logN) であることが魅力です。

まず配列 A, B, C を降順(値が大きい順)にソートします。すると A[i] + B[j] + C[k] が一番大きくなるのは A[0] + B[0] + C[0] です。

では二番目に大きいものは何でしょうか? A[1] + B[0] + C[0]、 A[0] + B[1] + C[0]、A[0] + B[0] + C[1] のどれかです。ここでは仮に二番目に大きなものを A[1] + B[0] + C[0] とします。すると三番目に大きなものは A[2] + B[0] + C[0]、A[1] + B[1] + C[0]、A[1] + B[0] + C[1] のどれかになります。

このようにどれかひとつの配列の添字をひとつだけ進めたものを優先度付きキューに格納していきます。そして格納したデータを取り出したらそのデータを元にそのどれかひとつの配列の添字をひとつだけ進めたものを優先度付きキューに格納していくという処理を繰り返します。これを K 回繰り返せばこの問題に正解することができます。

注意する点として優先度付きキューに一度格納した添字で生成されたデータを格納しないようにすることです。これは HashSet を使えば解決できそうです。

E – Set Meal

もう一題解いてみます。

E – Set Meal

食べ合わせが悪くならない組み合わせ以外でもっとも価格が高くなる定食の値段を求めよという問題です。

これも価格がもっとも高くなるものを求め、食べ合わせが悪い場合は配列の添字をそれぞれひとつずつズラしていけば解を得ることができます。

ソートするともとの添字がわからなくなってしまい、食べ合わせができなくなるので Pairクラスを定義してもともとの添字と値のペアを作ってます。あとは前述の問題と同じです。PriorityQueue を Dequeue することで価格が大きな順でメニューの組み合わせを得ることができるので、ここからもとの添字を取得して食べ合わせの確認をおこないます。問題ない組み合わせが見つかったらそれが解です。