BOJ_두 배열의 합_2143 (Java)

문제 링크

성능 요약

메모리: 112700 KB, 시간: 536 ms

분류

이분 탐색, 자료 구조, 해시를 사용한 집합과 맵, 누적 합

제출 일자

2025년 1월 13일 12:36:30

문제 설명

한 배열 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을 출력한다.

느낀점

• 처음엔 투포인터를 2중으로 사용해서 풀 수 있다고 생각했다. 하지만 원소에 음수가 포함될 수 있고, sum의 조건에 따라 구간을 조절하는 것이 아니라 모든 구간을 확인해야 하기 때문에 투포인터로는 확인이 안되는 부분이 있었다. (투포인터를 잘 모르는 걸수도,,)
• 1시간 30분정도 시도해보다가 설계가 잘못되었다는 걸 깨닫고 알고리즘 분류를 확인하니 누적합, 해시 등이 있었다.
• 알고리즘 분류를 보고서야 올바른 풀이를 떠올릴 수 있었다….
• 이분탐색도 있었는데, 이 문제에 어떻게 적용하는 방식인지는 모르겠다.
• 메모리 제한이 64MB라서 될까 싶었는데 가능했다. HashMap 자료구조를 좀 더 이해하기 위해 해싱 알고리즘에 대해 공부해야겠다는 생각이 들었다.

설계 : 40분

• 처음에는 투포인터를 통해 A에서 T보다 작은 부 배열의 합을 찾고, 나머지 값을 B에서 투포인터로 찾으려고 했는데 실패했다.
• 맞은 풀이
  • 입력된 A, B 배열을 누적합으로 저장한다.
  • 각 배열을 순회하며 각각의 Map에 구간합 값에 대한 경우의 수 카운트를 저장한다.
  • A 배열 구간합에 대한 Map의 Key를 순회하며 T - key 값을 B배열 구간합에 대한 Map에서 찾아 count끼리 곱해서 answer에 더한다.
  • ex) k + l =T일때, A에서 구간합으로 k가 나오는 경우가 3번이고, B에서 구간합으로 l이 나오는 경우가 2번이라고 하면 answer 에 3*2 = 6을 더한다.

코드(Java)

• 구현 시간: 100분

/**
 * Author: yngbao97, Yuk Yejin
 * Problem: 두 배열의 합_2143
 * Date: 2025.01.13
 */

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

public class Main {
	static BufferedReader br;
	static BufferedWriter bw;
	static StringTokenizer st;

	public static void main(String[] args) throws Exception {

		br = new BufferedReader(new InputStreamReader(System.in));
		bw = new BufferedWriter(new OutputStreamWriter(System.out));
		
		int T = Integer.parseInt(br.readLine());

		int[] prefixSumA = inputArr();
		int n = prefixSumA.length;
		int[] prefixSumB = inputArr();
		int m = prefixSumB.length;

		HashMap<Integer, Long> A = new HashMap<>();
		for (int i = 1; i < n; i++) {
			for (int j = i; j < n; j++) {
				int sum = prefixSumA[j] - prefixSumA[i-1];
				Long count = A.getOrDefault(sum, 0L);
				A.put(sum, count + 1);
			}
		}
		HashMap<Integer, Long> B = new HashMap<>();
		for (int i = 1; i < m; i++) {
			for (int j = i; j < m; j++) {
				int sum = prefixSumB[j] - prefixSumB[i-1];
				Long count = B.getOrDefault(sum, 0L);
				B.put(sum, count + 1);
			}
		}

		Long answer = 0L;
		for (int aSum : A.keySet()) {
			answer += A.get(aSum) * B.getOrDefault(T - aSum, 0L);
		}

		bw.write(String.valueOf(answer));
		bw.flush();
		bw.close();
		br.close();
	}

	private static int[] inputArr() throws IOException {
		int size = Integer.parseInt(br.readLine());
		int[] arr = new int[size+1];
		st = new StringTokenizer(br.readLine(), " ");
		for (int i = 1; i <= size; i++) arr[i] = arr[i-1] + Integer.parseInt(st.nextToken());

		return arr;
	}
}