[BOJ] 2143 두 배열의 합 - Java

🤔 문제

---

 

2143번: 두 배열의 합

 

😀 풀이

---

본 문제는 백준에 표기된 알고리즘 분류와 다르게 Map 자료구조를 사용하여 풀이했습니다.

 

문제에서 원하는 바는 주어진 두 배열의 부분합의 합이 문제에서 주어진 값(T)과 같아지는 경우가 몇 가지인지를 구하는 것 입니다. 

 

문제에서 주어진 입력의 크기가 1,000 임을 고려해보았을 때 브루트 포스로 풀이하게 되면 시간초과가 날 것이라 판단하여 아래와 같은 풀이를 사용했습니다.

1. A 배열의 부분합이 될 수 있는 값을 모두 구한다. + 각 값이 몇 번 나오는지 Map에 기록해둔다.
2. B 배열의 부분합이 될 수 있는 값을 모두 구한다. + 각 값이 몇 번 나오는지 Map에 기록해둔다.
3. 부분합 정보를 기록한 Map을 활용해 두 부분합의 합이 T 값이 될 수 있는  경우의 개수를 센다.

간략한 풀이 흐름은 위와 같은데 이제 좀 더 자세히 각 과정을 살펴보도록 하겠습니다.

 

1번과 2번 과정은 거의 동일하므로 1번 과정에 대해 설명해보겠습니다.

 

1번 과정의 경우 문제에서 준 예시처럼 A = {1, 3, 1, 2}으로 주어진 경우 부분합으로 나올 수 있는 경우는 총 (4 + 3 + 2 + 1) 가지 입니다. 아래와 같이 말이죠!

• 1, (1 + 3), (1 + 3 + 1), (1 + 3 + 1 + 2)
• 3, (3 + 1), (3 + 1 + 2)
• 1, (1 + 2)
• 2

이런 방식으로 부분합을 구하게 되면 이 정보를 Map에 담습니다. 그러면 Map에는 아래와 같이 정보가 담기게 될 것 입니다.

• subA = {1=2, 2=1, 3=2, 4=2, 5=1, 6=1, 7=1}

즉, Map의 key는 부분합, value는 그 부분합이 몇번 나오는지에 대한 정보가 담기게 되는 것이죠.

 

이 과정의 시간 복잡도는 O(N^2)이 됩니다. N이 1,000임을 감안한다면 문제가 되지 않습니다.

---

이제 위와 같은 과정을 통해 A와 B의 부분합 정보를 가진 Map을 통해 정답을 구하는 과정을 살펴보도록 하겠습니다.

위에서 언급한 3번 과정입니다.

 

문제에서 주어진 것 처럼 T가 5라면 부분합 정보를 가진 Map인 subA와 subB를 통해 T가 5가 될 수 있는 경우를 구해보겠습니다. 

 

subA의 첫 번째 key인 1의 경우를 살펴봅시다.

 

A의 부분합이 1이라면 B의 부분합이 4가 되어야 부분합의 합 T가 5가 될 수 있습니다.

 

이러한 점을 사용해 subB에서 key가 4인 경우의 value를 가져오고 가져온 두 value를 곱해주면 부분합의 합이 T, 즉, 5가 되는 경우의 수를 일부 구할 수 있습니다.

 

예시) subA.get(1) = 2, subB.get(4) = 1 인 경우 T = 5가 될 수 있는 경우는 이때 2가지를 발견했다는 것 입니다.  

 

위와 같은 방식으로 subA의 key를 모두 살펴보게 되면 T가 5가 되는 경우를 전부 구할 수 있게 됩니다.

---

💻 전체 풀이 코드는 아래와 같습니다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    static long T;
    static int N, M;
    static int[] A, B;
    static long answer = 0;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        T = Integer.parseInt(br.readLine());
        N = Integer.parseInt(br.readLine());
        
        A = new int[N];
        StringTokenizer st = new StringTokenizer(br.readLine(), " ");
        for (int i = 0; i < N; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }

        M = Integer.parseInt(br.readLine());
        B = new int[M];
        st = new StringTokenizer(br.readLine(), " ");
        for (int i = 0; i < M; i++) {
            B[i] = Integer.parseInt(st.nextToken());
        }

        Map<Long, Long> subA = new HashMap<>();
        Map<Long, Long> subB = new HashMap<>();

        // A의 부배열로 얻을 수 있는 합과 그 합의 개수 구하기
        for (int i = 0; i < N; i++) {
            long sum = 0;
            for (int j = i; j < N; j++) {
                sum += A[j];
                Long count = subA.getOrDefault(sum, 0L);
                subA.put(sum, count + 1);
            }
        }

        // B의 부배열로 얻을 수 있는 합과 그 합의 개수 구하기
        for (int i = 0; i < M; i++) {
            long sum = 0;
            for (int j = i; j < M; j++) {
                sum += B[j];
                Long count = subB.getOrDefault(sum, 0L);
                subB.put(sum, count + 1);
            }
        }

        for (Map.Entry<Long, Long> entry : subA.entrySet()) {
            Long aKey = entry.getKey();
            Long aValue = entry.getValue();
            long bKey = T - aKey;
            Long bValue = subB.getOrDefault(bKey, 0L);
            answer += (aValue * bValue);
        }

        System.out.println(answer);
    }
}