概要
エンジニア職での就職・転職では会社によってはコーディングテストというものがあります。
新卒でも中途でも入社したい会社というのは誰でもいい環境を求めるのではないでしょうか。
そこで、コーディングテストで振り落とされないためにもLeetCodeというテスト対策ができるサイトでアルゴリズム力を鍛えたいと思います。
問題
373. Find K Pairs with Smallest Sums
昇順に並べられた2つの整数配列nums1、nums2と、整数kが与えられる。
最初の配列から1つの要素、2番目の配列から1つの要素からなるペア(u, v)を定義しなさい。
和が最も小さい方からk番目のペア(u1, v1), (u2, v2), …, (uk, vk)を返せ.
制約
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1
とnums2
は両方とも昇順にソートされている1 <= k <= 104
解説
考え方
まず初めにnum1から値をk個分取り出しそれぞれをnum2の最初の値とペアにしたものをmin_heapに入れていきます。num1はすでにソートされているので、自動的に合計数の昇順になるはずです。
もし、num1の個数がkよりも少なかった場合はnum1の数だけペアにしましょう。
次に、min_heapの一番小さい値を取り出し、出力用のリストに追加します。・・・①
その際、追加する際に使用した「num1の値」と「num2の次の値」をペアにしてmin_heapに追加します。・・・②
(num2の最後の値を出力用リストに追加した場合は次の値をペアにできないので、スキップします。)
後は①と②を繰り返して出力用リストにデータを追加しましょう。
データがkの数まで用意できるか、min_heapが空になって追加できなくなれば出力用リストを返して終了です。
PriorityQueueでの解き方
まずはnum1とnum2から取り出したペアのインデックスを格納するクラスを作成する。
配列でも可能ですが、読みやすいようにオブジェクト化しています。
public class IndexPair
{
public readonly int Index1;
public readonly int Index2;
public IndexPair(int index1, int index2)
{
Index1 = index1;
Index2 = index2;
}
}
IndexPair
クラスを使用してPriorityQueue
を作成します。
その後num1とnum2の一番小さいデータを順にIndexPair
を作成してPriorityQueue
に追加します。
この際、ペアの和の昇順としたいので第2引数の優先度にはペアの和を渡します。
// インデックスをペアにしたデータと合計を優先度とした初期値を作成する
var priorityQueue = new PriorityQueue<IndexPair, int>();
// num1の数もしくはkの数までデータを追加する
for (int i = 0; i < nums1.Length && i < k; i++)
{
priorityQueue.Enqueue(new IndexPair(i, 0), nums1[i] + nums2[0]);
}
次に出力用リストを用意します。KSmallestPairs
メソッドではILIST<ILIST<int>>
を戻り値としているのでnew List<ILIST<int>>
でインスタンスを生成します。
// 出力用リストを作成
IList<IList<int>> result = new List<IList<int>>();
出力用リストがk個分用意できるか、PriorityQueue
が空になるまでデータを追加するので、while文の条件は「出力用リストのデータ数がkよりも少ない かつ 優先度付きキューにデータがある」となります。
キューからデータを取り出して、出力用リストに追加します。
優先度付きキューで合計の昇順にソートされているので自然と一番合計が小さいペアが取り出せます。
データを取り出したら次のペアを検索して優先度キューに追加してください。
後は繰り返せばデータが揃って終了です。
// 出力用リストがk個分用意できるまで かつ 優先度付きキューにデータがある場合繰り返す
while (result.Count < k && priorityQueue.Count > 0)
{
// 優先度付きキューからインデックスを取り出したインデックスを使用してペアを出力用データとして追加
var dequeued = priorityQueue.Dequeue();
int index1 = dequeued.Index1;
int index2 = dequeued.Index2;
result.Add(new List<int>() { nums1[index1], nums2[index2] });
// num2で使用したインデックスを+1して次のペアを優先度キューに追加する
index2++;
if (index2 < nums2.Length)
{
priorityQueue.Enqueue(
new IndexPair(index1, index2),
nums1[index1] + nums2[index2]);
}
}
ソースコード
public class Solution {
public class IndexPair
{
public readonly int Index1;
public readonly int Index2;
public IndexPair(int index1, int index2)
{
Index1 = index1;
Index2 = index2;
}
}
public IList<IList<int>> KSmallestPairs(int[] nums1, int[] nums2, int k)
{
// インデックスをペアにしたデータと合計を優先度とした初期値を作成する
var priorityQueue = new PriorityQueue<IndexPair, int>();
for (int i = 0; i < nums1.Length && i < k; i++)
{
priorityQueue.Enqueue(new IndexPair(i, 0), nums1[i] + nums2[0]);
}
// 出力用リストを作成
IList<IList<int>> result = new List<IList<int>>();
// 出力用リストがk個分用意できるまで かつ 優先度付きキューにデータがある場合繰り返す
while (result.Count < k && priorityQueue.Count > 0)
{
// 優先度付きキューからインデックスを取り出したインデックスを使用してペアを出力用データとして追加
var dequeued = priorityQueue.Dequeue();
int index1 = dequeued.Index1;
int index2 = dequeued.Index2;
result.Add(new List<int>() { nums1[index1], nums2[index2] });
// num2で使用したインデックスを+1して次のペアを優先度キューに追加する
index2++;
if (index2 < nums2.Length)
{
priorityQueue.Enqueue(
new IndexPair(index1, index2),
nums1[index1] + nums2[index2]);
}
}
return result;
}
}