문제

https://www.acmicpc.net/problem/2143

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력
첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력
첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

예제 입력 1
5
4
1 3 1 2
3
1 3 2
예제 출력 1
7

접근

완전 해시맵 문제라고 생각이 들긴 했는데, 메모리 제한이 64MB라 괜히 겁을 먹고... 맵을 안 쓰는 방향으로 구현을 해보려고 했었다. 이분 탐색이나 투 포인터로 가능할 것 같았는데, 이분 탐색으로 해보니 계속 틀렸대서... 그냥 해시맵으로 해봤는데 맞았다..? 음....
A의 부배열을 관리하는 해시맵과 B의 부배열을 관리하는 해시맵을 각각 생성한다. 그리고 해시맵의 key 값으로는 가능한 부 배열의 합을, value 값으로는 해당 부 배열 합이 몇 번 등장하는지 횟수를 저장해주면 된다.

구현

1차 시도 (이분 탐색, 틀린 풀이)

import java.io.*;
import java.util.*;

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

        List<long[]> sumOfA = new ArrayList<>();
        for (int start = 0; start < n; start++) {
            sumOfA.add(new long[] {A[start], 0});
            for (int j = start + 1; j < n; j++) {
                sumOfA.add(new long[] {sumOfA.get(sumOfA.size() - 1)[0] + A[j], 0});
            }
        }
        Collections.sort(sumOfA, (o1, o2) -> {return (int)(o1[0] - o2[0]);});

        List<Integer> sumOfB = new ArrayList<>();
        for (int start = 0; start < m; start++) {
            sumOfB.add(B[start]);
            for (int j = start + 1; j < m; j++) {
                sumOfB.add(sumOfB.get(sumOfB.size() - 1) + B[j]);
            }
        }
        Collections.sort(sumOfB);

        long answer = 0;
        for (int k = 0; k < sumOfA.size(); k++) {
            if (k > 1 && sumOfA.get(k)[0] == sumOfA.get(k - 1)[0]) {
                answer += sumOfA.get(k - 1)[1];
                sumOfA.get(k)[1] = sumOfA.get(k - 1)[1];
                continue;
            }
            int location = Collections.binarySearch(sumOfB, (int)(T - sumOfA.get(k)[0]));
            if (location < 0) {
                continue;
            }
            sumOfA.get(k)[1]++;
            int before = location - 1;
            while (before >= 0 && sumOfB.get(before) == sumOfB.get(location)) {
                sumOfA.get(k)[1]++;
                before--;
            }
            int after = location + 1;
            while (after < sumOfB.size() && sumOfB.get(after) == sumOfB.get(location)) {
                sumOfA.get(k)[1]++;
                after++;
            }
            answer += sumOfA.get(k)[1];
        }

        System.out.println(answer);
    }
}

약간 하드 코딩 느낌이 있지만 논리상 틀린 건 없어 보이는데 왜 틀린 건지 모르겠다...
나는 이분 탐색을 한 번 수행한 후에 그 찾은 위치의 양 옆을 확인하며 같은 수가 있는지 확인했는데, 이분 탐색을 쓴 정답 풀이를 보니 lower bound와 upper bound를 각각 찾아서 그 사이 개수를 세는 식으로 하더라.
후자가 더 멋들어져 보이긴 하는데... 아직도 뭐가 틀린 건지는 잘 모르겠다.

2차 시도(해시맵, 정답 풀이)

import java.io.*;
import java.util.*;

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

        Map<Integer, Long> sumOfA = new HashMap<>();
        for (int start = 0; start < n; start++) {
            int sum = 0;
            for (int end = start; end < n; end++) {
                sum += A[end];
                sumOfA.put(sum, sumOfA.getOrDefault(sum, 0L) + 1);
            }
        }

        Map<Integer, Long> sumOfB = new HashMap<>();
        for (int start = 0; start < m; start++) {
            int sum = 0;
            for (int end = start; end < m; end++) {
                sum += B[end];
                sumOfB.put(sum, sumOfB.getOrDefault(sum, 0L) + 1);
            }
        }

        long answer = 0;
        for (Integer subA : sumOfA.keySet()) {
            if (sumOfB.keySet().contains(T - subA)) {
                answer += (sumOfA.get(subA) * sumOfB.get(T - subA));
            }
        }

        System.out.println(answer);
    }
}

이 문제는 솔직히 해시맵으로 풀면 남는 게 별로 없는 문제 같다...
다음에 투 포인터로 다시 풀어봐야겠으....