[LeetCode] 373. Find K Pairs with Smallest Sums

Find K Pairs with Smallest Sums - LeetCode

오름차순 정렬된 두 개의 정수 배열 nums1과 nums2와 정수 k가 주어집니다.
첫 번째 배열의 원소 1개와 두 번째 배열의 원소 1개로 구성된 쌍(u, v)을 정의합니다.
합이 가장 작은 (u1, v1), (u2, v2), ..., (uk, vk)의 k쌍을 반환합니다.

 

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

 

문제접근

배열1, 배열2가 서로 조합된 배열이 필요하고

조합된 요소들의 합이 작은순서대로 반환해야 합니다.

 

따라서 DFS를 통해 두 배열의 조합을 구하고

정렬하여 반환하면 되겠다고 생각했습니다.

 

DFS함수는 nums1의 원소를 하나씩 temp에 넣어서

nums2와 조합을 만듭니다.

temp.size == 2라면 temp의 합을 add해서 combination에 담아주었습니다.

 

그 후 Stream을 통해 합을 기준으로 정렬하고

합의 원소를 빼주고

limit을 k번까지만 하여 반환해 주었습니다.

 

class Solution {
    ArrayList<ArrayList<Integer>> combination = new ArrayList<>();

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

        boolean[] visited = new boolean[nums2.length];
        ArrayList<Integer> temp = new ArrayList<>();
        for (int i = 0; i < nums1.length; i++) {
            temp.add(nums1[i]);
            dfs(nums2,temp, visited);
            temp.remove(temp.size()-1);
        }
        final List<List<Integer>> result = combination.stream()
                .sorted((a1, a2) -> a1.get(2) - a2.get(2))
                .map(a1 -> Arrays.asList(a1.get(0), a1.get(1)))
                .limit(k)
                .collect(Collectors.toList());

        return result;
    }

    private void dfs(int[] nums, ArrayList<Integer> temp, boolean[] visited) {
        if (temp.size() == 2) {
            temp.add((temp.get(0) + temp.get(1)));
            combination.add(new ArrayList<>(temp));
            temp.remove(2);
            return;
        }

        for (int i = 0; i < nums.length; i++) {
            if (visited[i]) continue;
            visited[i] = true;
            temp.add(nums[i]);
            dfs(nums, temp, visited);
            visited[i] = false;
            temp.remove(temp.size()-1);
        }
    }
}

하지만, 메모리 초과로 통과하지 못했습니다.

 

처음에는 ArraysList가 메모리를 많이 잡아먹었나? 라고 생각해서

int[]로 바꿔보았지만 여전히 메모리가 초가되었습니다.

 

INT배열 풀이

class Solution {
     ArrayList<int[]> permutation = new ArrayList<>();

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

        boolean[] visited = new boolean[nums2.length];
        int[] temp = new int[3];
        for (int i = 0; i < nums1.length; i++) {
            temp[0] = nums1[i];
            dfs(nums2,temp, 1, visited);
        }
        final List<List<Integer>> result = permutation.stream()
                .sorted((a1, a2) -> a1[2] - a2[2])
                .map(a1 -> Arrays.asList(a1[0], a1[1]))
                .limit(k)
                .collect(Collectors.toList());

        return result;
    }

    private void dfs(int[] nums, int[] temp, int depth, boolean[] visited) {
        if (depth == 2) {
            temp[2] = temp[0] + temp[1];
            permutation.add(Arrays.copyOf(temp, temp.length));
//            System.out.println(Arrays.toString(temp));
            return;
        }

        for (int i = 0; i < nums.length; i++) {
            if (visited[i]) continue;
            visited[i] = true;
            temp[depth] = nums[i];
            dfs(nums, temp, depth+1, visited);
            visited[i] = false;
            temp[depth] = 0;
        }
    }
}

 

 

nums1, nums2의 범위가 10의 5승이여서 초과가 된것입니다.

 

다른사람의 풀이

PriroytyQueue를 사용하였습니다.

 

우선 PriroytyQueue를 sum기준 오름차순 정렬합니다.

그리고 반복문을 돌면서 PriroytyQueue의 size가 k개와 같을때까지 offer해주고

PriroytyQueue의 size가 k개라면, 현재 PriroytyQueue의 sum과 대기중인 sum을 비교하여

PriroytyQueue의 sum이 더크다면 다시 PriroytyQueue에 offer하는 것입니다.

 

즉, PriroytyQueue의 최대 크기를 K개로 두어서

정답을 확인할 수 있습니다.

public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> (b[0] + b[1])-(a[0]+a[1]));

        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            for (int j = 0; j < Math.min(nums2.length, k); j++) {
                int[] cur={nums1[i], nums2[j]};
                if (pq.size() == k) {
                    int[] pk=pq.peek();
                    if(pk[0]+pk[1]>cur[0]+cur[1]){
                        pq.poll();
                        pq.offer(cur);
                    }
                    else break;
                }
                else pq.offer(cur);
            }
        }

        List<List<Integer>> result = new ArrayList<>();
        while (!pq.isEmpty()) {
            int[] pair = pq.poll();
            result.add(Arrays.asList(pair[0], pair[1]));
        }

        return result;
    }

 

느낀점

PriroytyQueue를 더 잘 활용해야겠다는 느낌이 들었습니다.

PriroytyQueue를 사용해서 문제를 풀어야겠다는 생각은 하지 못했는데

다양한 풀이법을 생각해봐야겠습니다.