BOJ_2143_두 배열의 합 (Java)

[Gold III] 두 배열의 합 - 2143

문제 링크

성능 요약

메모리: 73020 KB, 시간: 340 ms

분류

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

제출 일자

2024년 9월 17일 00:35:10

문제 설명

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

문제 풀이

누적합으로 문제를 풀어보았다. 누적합 배열 a, b, 에서 i까지의 값은 1~i까지의 합을 나타낸다. 이후 HashMap 자료구조를 통해 일치 조건을 탐색했다.

Key값으로 합, value값으로 횟수를 넣었다. 특정 구간의 i~j값은 배열에서 arr[j] - arr[i]로 구할 수 있고 이때의 값을 Key로 하여 나오는 횟수를 value로 넣었다. map의 getOrDefault함수는 값이 있으면 그 키 값의 값을 사용, 없으면 0을 사용하는 메서드다. 매우 유용하니 Map자료구조를 사용하면 꼭 사용해보자.

이렇게 한 배열을 순회하였으면 나머지 배열을 순회하며 모든 구간합마다 T-sum를 key값으로 가지는 것이 있으면 그 횟수만큼 더해주는 식으로 풀었다.

코드

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.StringTokenizer;

public class Main {
	static BufferedReader br;
	static StringTokenizer st;
	static int a[], b[];
	static HashMap<Long, Integer> map;
	public static void main(String[] args) throws IOException {
//		br = new BufferedReader(new InputStreamReader(new FileInputStream("input.txt")));
		br = new BufferedReader(new InputStreamReader(System.in));
		
		int T = Integer.parseInt(br.readLine());
		
		int N = Integer.parseInt(br.readLine());
		a = new int[N+1];
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++) {
			a[i] = a[i-1] + Integer.parseInt(st.nextToken());
		}
		
		int M = Integer.parseInt(br.readLine());
		b = new int[M+1];
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=M; i++) {
			b[i] = b[i-1] + Integer.parseInt(st.nextToken());
		}
		
		map = new HashMap<Long, Integer>();
		long ans = 0;
		
		for(int i=0; i<=M; i++) {
			for(int j=i+1; j<=M; j++) {
				long sum = b[j] - b[i];
				map.put(sum, map.getOrDefault(sum, 0)+1);
			}
		}
		
		for(int i=0; i<=N; i++) {
			for(int j=i+1; j<=N; j++) {
				long sum = a[j] - a[i];
				ans += map.getOrDefault(T-sum, 0);
			}
		}
		
		System.out.println(ans);
		
	}
}