[Java] 백준 2143번: 두 배열의 합

[문제 바로 가기] - 두 배열의 합

유형 : 누적합 + 투 포인터

💡 접근 방식

부배열이란 해당 배열에서 연속되는 부분 배열이므로 배열의 누적합을 구한다.
예를 들면 A = [1, 3, 1, 2]일 때

1	3	1	2
1	4	5	7
	3	4	6
		1	3
			2

위의 표와 같이 구해진다. 따라서 A와 B의 각 누적합을 구해서 listA와 listB에 넣어주고 Collecions.sort로 정렬해주었다.

listA와 listB에서 하나씩 값을 가져와 합이 T가 되는 경우의 수를 구하는데, 처음에는 이중 for문을 돌면서 하나씩 구했다.

그랬더니 당연히 시간 초과가 났다!

시간 단축을 위해 합이 T가 되는 경우의 수를 구할 때는 투 포인터를 사용했다.

1. 합이 T와 같을 때 -> answer++
2. 합이 T보다 작을 때 -> left++
3. 합이 T보다 클 때 -> right++

1번에서 listA와 listB에 같은 수가 존재할 수 있으므로 aCount와 bCount를 센 후 곱한 것을 answer에 더해주어야 한다.

★ 주의할 점!

• aCount와 bCount가 int 타입을 벗어날 수 있으므로 long 타입으로 선언해야 함
• List의 원소를 비교할 땐 equals 함수를 사용해야 함
  - listA.get(left) == listA.get(left + 1)은 값 비교가 아니라 객체 참조를 비교하는 것

---

풀이

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.StringTokenizer;

/**
 * 백준 2143번 두 배열의 합 
 * - 누적합 + 투 포인터
 * - 타입, List 값 비교 주의
 */

public 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<Integer> listA = new ArrayList<>();
		List<Integer> listB = new ArrayList<>();
		
		for (int i = 0; i < N; i++) {
			int sum = 0;
			for (int j = i; j < N; j++) {
				sum += A[j];
				listA.add(sum);
			}
		}
		
		for (int i = 0; i < M; i++) {
			int sum = 0;
			for (int j = i; j < M; j++) {
				sum += B[j];
				listB.add(sum);
			}
		}
		
		Collections.sort(listA);
		Collections.sort(listB);
		
		long answer = 0;
		
		int left = 0;
		int right = listB.size() - 1;
		
		while (left < listA.size() && right >= 0) {
			int sum = listA.get(left) + listB.get(right);
			
			if (sum == T) {
				long aCount = 1;
				long bCount = 1;
				
				while (left + 1 < listA.size() && listA.get(left).equals(listA.get(left + 1))) {
					left++;
					aCount++;
				}
				
				while (right - 1 >= 0 && listB.get(right).equals(listB.get(right - 1))) {
					right--;
					bCount++;
				}
				
				answer += aCount * bCount;
						
				right--;
				left++;
			} else if (sum < T) left++;
			else right--;
		}
		
		System.out.println(answer);
	}
}