Daily LeetCode Challenge - 373. Find K Pairs with Smallest Sums

Problem From.
https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

오늘 문제는 nums1 과 nums2 가 주어질때, 각각의 원소들을 더해서 작은 값이 되는 수들을 짝지어서 k 개수만큼 반환하는 문제였다.

이 문제는 priority queue 를 이용해서 풀 수 있었는데, priority queue 에 각각의 원소의 합들을 인덱스를 하나씩 추가해주면서 넣어주고, 다시 queue 가 빌때까지 꺼내오면서 list 에 추가해주면 되었다. priority queue 가 각 queue 의 값들을 오름차순으로 가지고 있기때문에, queue 에서 꺼내와서 반환해주면 되었다.

class Solution {
    fun kSmallestPairs(nums1: IntArray, nums2: IntArray, k: Int): List<List<Int>> {
        val n = nums1.size
        val m = nums2.size
        val visited = mutableSetOf<String>()
        var queue = PriorityQueue<IntArray>() { arr1, arr2 ->
            (nums1[arr1[0]] + nums2[arr1[1]]).compareTo(nums1[arr2[0]] + nums2[arr2[1]])
        }
        queue.add(intArrayOf(0, 0))

        val answer = mutableListOf<List<Int>>()
        repeat(k){
            if(queue.isEmpty()) return answer
            val (i, j) = queue.poll()
            answer.add(listOf(nums1[i], nums2[j]))

            val k1 = "${i + 1}, $j"
            if(i + 1 < n && !visited.contains(k1)){
                queue.add(intArrayOf(i + 1, j))
                visited.add(k1)
            }

            val k2 = "${i}, ${j + 1}"
            if(j + 1 < m && !visited.contains(k2)){
                queue.add(intArrayOf(i, j + 1))
                visited.add(k2)
            }
        }

        return answer
    }
}