diff --git a/373. Find K Pairs with Smallest Sums.md b/373. Find K Pairs with Smallest Sums.md new file mode 100644 index 0000000..456f789 --- /dev/null +++ b/373. Find K Pairs with Smallest Sums.md @@ -0,0 +1,262 @@ +# step 1 +nums1, nums2のインデックスを一つずつ動かす-> wrong answer +全パターン試せていなかった。 +`nums1 = [1, 2, 1000], nums2 = [1, 100], k = 3`でうまくいかなくなる。 +```python +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + if len(nums1) * len(nums2) < k: + raise ValueError( + "k is greater than the number of all possible pairs: " + f"k = {k}, nums1 = {nums1}, nums2 = {nums2}" + ) + + k_smallest_pairs = [] + index1 = 0 + index2 = 0 + while len(k_smallest_pairs) < k: + k_smallest_pairs.append([nums1[index1], nums2[index2]]) + + if index1 + 1 == len(nums1): + index2 += 1 + continue + if index2 + 1 == len(nums2): + index1 += 1 + continue + + if nums1[index1 + 1] <= nums2[index2 + 1]: + index1 += 1 + else: + index2 += 1 + return k_smallest_pairs +``` + +leetcodeのeditorialを見て実装. +```python +import heapq + + +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + k_smallest_pairs = [] + visited = set() + next_pair = [] + + heapq.heappush(next_pair, (nums1[0] + nums2[0], 0, 0)) + while len(k_smallest_pairs) < k: + _, index1, index2 = heapq.heappop(next_pair) + + if (index1, index2) in visited: + continue + + k_smallest_pairs.append([nums1[index1], nums2[index2]]) + visited.add((index1, index2)) + if index1+1 < len(nums1) and (index1+1, index2) not in visited: + heapq.heappush( + next_pair, + (nums1[index1+1] + nums2[index2], index1+1, index2) + ) + if index2+1 < len(nums2) and (index1, index2+1) not in visited: + heapq.heappush( + next_pair, + (nums1[index1] + nums2[index2+1], index1, index2+1) + ) + return k_smallest_pairs +``` + +全てのパターンを試して、小さい順にheapにkだけ詰めて管理することを考えていたが、kが大きくなった時にTLEすると感じた。 + +回答の方は次に探索するインデックスの方をheapで管理していた。グラフ探索みたいだった。 +leetcodeの解答を少し書き換えて、visited.add()を呼び出す場所を自分の感覚に合う様にした。後述したが、これは変数名を`explored`とした方が適切かも。 +その結果として、heapqに同じ探索先が2度入りうることになった(ex. `nums1 = [1, 2, 3], nums2 = [1, 2]`で、next_pairに(4, 1, 1)が2つ入ることになる)。 + +# step 2 +- https://github.com/tarinaihitori/leetcode/pull/10/files + - `num1_i`, `num2_i`と書かれている箇所があり、何度か見間違えた。似た様な意味の変数だと先頭か末尾が違う方が好み。 + - ほぼ同じ解法で`visited`を定義していなかった。 + - 重複なく追加するために、以下の様に書いていた。 + ```python + if j == 0 and i + 1 < len(nums1): + heapq.heappush(sum_with_each_index, (nums1[i + 1] + nums2[0], i + 1, 0)) + if j + 1 < len(nums2): + heapq.heappush(sum_with_each_index, (nums1[i] + nums2[j + 1], i, j + 1)) + ``` + (i, j)にたどり着くためのルートが一本に定まる様にしてある。 + + 個人的には、自分の書いたもののほうが読みやすかった。自分の書いたものは対称的かつ、dijkstraの延長上のものと理解できること、こちらのコードは微妙に対称性が崩れており、その理由を考えるのに時間がかかることが原因だと思う。 + - `next_pair`変数の代わりに、`sum_with_each_index`という変数名を使っていた。 +- https://github.com/goto-untrapped/Arai60/pull/56 + - heappushする箇所を関数化するのもあり + - leetcodeの解答同様、heapにpushした時点でvisitedとしていた。 + - step1で書いたコードではheappopしたタイミングをvisitedとしていたが、exploredのほうが適切かも。 +- https://github.com/colorbox/leetcode/pull/25/files + - 二分探索を使った解法があった。numの配列が非減少->二分探索と思考を進めるとこういう解法も出てくるのかと面白かった。 + +```python +import heapq + + +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + visited = set() + next_pair_candidates = [] + k_smallest_pairs = [] + + heapq.heappush(next_pair_candidates, (nums1[0]+nums2[0], 0, 0)) + visited.add((0, 0)) + while len(k_smallest_pairs) < k: + _, index1, index2 = heapq.heappop(next_pair_candidates) + k_smallest_pairs.append([nums1[index1], nums2[index2]]) + + if index1 + 1 < len(nums1) and (index1 + 1, index2) not in visited: + heapq.heappush( + next_pair_candidates, + (nums1[index1 + 1] + nums2[index2], index1 + 1, index2) + ) + visited.add((index1 + 1, index2)) + if index2 + 1 < len(nums2) and (index1, index2 + 1) not in visited: + heapq.heappush( + next_pair_candidates, + (nums1[index1] + nums2[index2 + 1], index1, index2 + 1) + ) + visited.add((index1, index2 + 1)) + return k_smallest_pairs +``` + +# step 3 + +```python +import heapq + + +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + visited = set() + next_pair_candidates = [] + k_smallest_pairs = [] + + heapq.heappush(next_pair_candidates, (nums1[0] + nums2[0], 0, 0)) + visited.add((0, 0)) + while len(k_smallest_pairs) < k: + _, index1, index2 = heapq.heappop(next_pair_candidates) + k_smallest_pairs.append([nums1[index1], nums2[index2]]) + + if index1 + 1 < len(nums1) and (index1 + 1, index2) not in visited: + visited.add((index1 + 1, index2)) + heapq.heappush( + next_pair_candidates, + (nums1[index1 + 1] + nums2[index2], index1 + 1, index2) + ) + if index2 + 1 < len(nums2) and (index1, index2 + 1) not in visited: + visited.add((index1, index2 + 1)) + heapq.heappush( + next_pair_candidates, + (nums1[index1] + nums2[index2 + 1], index1, index2 + 1) + ) + return k_smallest_pairs +``` + +# step 4 +コメントまとめ +- `+`の両端にスペースを開ける +- visitedはグラフ探索っぽい。今回の文脈ではprocessedのほうが適切 +- 次のノードの決定法で自分の見ていなかった回答がいくつかあった + - index1がある値のときにindex2をどこまで探索したかと、それと逆に、index2がある値のときにindex1をどこまで探索したのかを記録する方法：https://github.com/TORUS0818/leetcode/commit/09636a11c0c084c245f82bb029f198c2d0966455#diff-12341a70069caa7ec44aba610c42a0123605c2680c11a88728158764788bead7R226 + - step2の一番最初のものとほとんど一緒だが、最初にheapqにindex2=0となる全ての場合をpushしてから探索を始めるもの: https://github.com/seal-azarashi/leetcode/pull/10/files/051cdb9294e7cef521f43d05c394e35c40a3d52a + - step2の初めのものよりは読みやすかった。次のノードをpushする際にindex2がひたすらインクリメントされていくだけなので、(i, j)に至るルートは同じだが、わかりやすい。 + + +共通処理の関数化をした解法 +- 初めはnums1, nums2も引数にしていたが、引数の数が大きくなりすぎるのもと思い取り除いた。 + - 引数にない、書き換え可能な値を参照することが気持ち悪い + - pythonだと関数引数にconst指定できないのが少し不便 +- inner functionと例外処理の順番をどうしようか迷った。例外処理してから解法に関係のあるコードを書こうと思い、そう書いた。 +```python +import heapq + + +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + if not 1 <= k <= len(nums1) * len(nums2): + raise ValueError( + 'kSmallestPairs(): k is out of range :' + f"nums1 = {nums1}, nums2 = {nums2}, k = {k}" + ) + + def add_next_candidate( + next_pairs: List[Tuple[int, Tuple[int, int]]], + index1: int, + index2: int, + processed: Set[Tuple[int, int]] + ) -> None: + if ( + index1 < len(nums1) + and index2 < len(nums2) + and (index1, index2) not in processed + ): + heapq.heappush( + next_pairs, + (nums1[index1] + nums2[index2], (index1, index2)) + ) + + k_smallest_pairs = [] + next_pairs = [] + processed = set() + + heapq.heappush(next_pairs, (nums1[0] + nums2[0], (0, 0))) + while len(k_smallest_pairs) < k: + _, (index1, index2) = heapq.heappop(next_pairs) + if (index1, index2) in processed: + continue + + k_smallest_pairs.append([nums1[index1], nums2[index2]]) + processed.add((index1, index2)) + + add_next_candidate(next_pairs, index1 + 1, index2, processed) + add_next_candidate(next_pairs, index1, index2 + 1, processed) + return k_smallest_pairs +``` + +mergeしていくやり方 + +はじめに(i, 0)のペアを作る。jを1から順にnums2.lengthまで動かしながら、あるjにおける(i, j)のリストを作り(これは非減少)、これまでの最小値のとなるkペアにたいしてmergeする。 +- time complexity: O(k^2) +- space complexity: O(k) +```python3 +import heapq + + +class Solution: + def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: + k_smallest_pairs = [] + for i in range(len(nums1)): + if len(k_smallest_pairs) == k: + break + k_smallest_pairs.append((nums1[i] + nums2[0], nums1[i], nums2[0])) + + for index2 in range(1, len(nums2)): + if index2 == k: + break + if ( + len(k_smallest_pairs) == k + and k_smallest_pairs[-1][0] <= nums1[0] + nums2[index2] + ): + break + smaller_pairs = [] + for index1 in range(len(nums1)): + if ( + len(k_smallest_pairs) == k + and k_smallest_pairs[-1][0] <= nums1[index1] + nums2[index2] + ): + break + if len(smaller_pairs) == k: + break + smaller_pairs.append(( + nums1[index1] + nums2[index2], + nums1[index1], + nums2[index2] + )) + k_smallest_pairs = list(heapq.merge(k_smallest_pairs, smaller_pairs))[:k] + + return [[num1, num2] for _, num1, num2 in k_smallest_pairs] +``` \ No newline at end of file