[leetcode #454] 4Sum II

Problem

Given four integer arrays nums1, nums2, nums3, and nums4 all of length n, return the number of tuples (i, j, k, l) such that:

・ 0 <= i, j, k, l < n
・ nums1[i] + nums2[j] + nums3[k] + nums4[l] == 0

Example 1:

Input: nums1 = [1,2], nums2 = [-2,-1], nums3 = [-1,2], nums4 = [0,2]
Output: 2
Explanation:
The two tuples are:
1. (0, 0, 0, 1) -> nums1[0] + nums2[0] + nums3[0] + nums4[1] = 1 + (-2) + (-1) + 2 = 0
2. (1, 1, 0, 0) -> nums1[1] + nums2[1] + nums3[0] + nums4[0] = 2 + (-1) + (-1) + 0 = 0

Example 2:

Input: nums1 = [0], nums2 = [0], nums3 = [0], nums4 = [0]
Output: 1

Constraints:

・ n == nums1.length
・ n == nums2.length
・ n == nums3.length
・ n == nums4.length
・ 1 <= n <= 200
・ -2²⁸ <= nums1[i], nums2[i], nums3[i], nums4[i] <= 2²⁸

Idea

어렵게 생각할 필요가 없는 문제다. 주어진 네 개의 array를 두 그룹으로 나누어 각 그룹에서 하나씩 뽑은 수로 만든 합을 map에 저장해서 풀면 된다. 이 때 map의 key는 합, value는 합이 가능한 수다.

각 group의 모든 값을 합과 가능한 경우의 수로 분리한 뒤, 한 그룹의 합을 key로 하여 다른 그룹의 key와 합이 0이 되는 경우가 있는지 확인한다. 만약 0이라면 경우의 수를 서로 곱해 결과값에 더하면 된다.

모든 key를 탐색한 뒤 결과값을 리턴하면 된다.

Time Complexity: O(n²)
Space Complexity: O(n²)

Solution

class Solution {
    public int fourSumCount(int[] nums1, int[] nums2, int[] nums3, int[] nums4) {
        Map<Integer, Integer> first = new HashMap<>();
        Map<Integer, Integer> second = new HashMap<>();

        for (int num1 : nums1) {
            for (int num2 : nums2) {
                int sum = num1 + num2;
                int cnt = first.getOrDefault(sum, 0);
                first.put(sum, cnt+1);
            }
        }

        for (int num3 : nums3) {
            for (int num4 : nums4) {
                int sum = num3 + num4;
                int cnt = second.getOrDefault(sum, 0);
                second.put(sum, cnt+1);
            }
        }

        int res = 0;
        for (int key : first.keySet()) {
            if (!second.containsKey(-key)) {
                continue;
            }
            res += first.get(key) * second.get(-key);
        }

        return res;
    }
}

Reference

https://leetcode.com/problems/4sum-ii/