[BOJ] 백준 2143번 두 배열의 합 (Java)

#2143 두 배열의 합

난이도 : 골드 3

유형 : 이분 탐색/ 투포인터

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

 

문제 풀이  

각 배열의 부 배열의 합 모든 경우를 구한 다음 T에 맞는 원소들을 매칭해주면 된다. N개의 원소를 갖는 배열의 부 배열의 크기는 n*(n-1)/2이므로 두 배열의 부배열을 구하는 데 걸리는 시간 O(N^2)이다. 그런 다음 두 부 배열 원소의 합이 T가 되는 값들을 매칭시켜주면 된다. 매칭시켜주는 방법으로는 투 포인터와 이진탐색 방법이 있다. 투 포인터로 푸는 것이 더 빠르고 편리하다. 그래도 연습이니 두 케이스 풀이를 모두 진행해봤다.

 

부 배열은 각 원소들의 구간 합을 모두 구해주면 된다. 

int aSize = n*(n+1)/2;
int bSize = m*(m+1)/2;
long[] aSum = new long[aSize];
int idx=0;
for(int i=0; i<n; i++) {
	int av=0;
	for(int j=i; j<n; j++) {
		av+=a[j];
		aSum[idx++] = av;
	}
}

long[] bSum = new long[bSize];
idx=0;
for(int i=0; i<m; i++) {
	int bv = 0;
	for(int j=i; j<m; j++) {
		bv += b[j];
		bSum[idx++] = bv;
	}
}
Arrays.sort(aSum);
Arrays.sort(bSum);

 

이제 위에서 구한 두 개의 부배열의 원소들 중에서 합이 T가 되는 조합을 구해주면 된다.

투 포인터

투 포인터 방법으로는 두 배열의 합이 T가 되는 구간을 한 칸씩 이동하면서 탐색해주면 된다. 중복되는 원소도 처리해주기 때문에 빠르게 탐색이 가능하다.

1. aSum은 0번째 부터 탐색을 시작한다. bSum은 bSize-1부터 탐색을 시작한다.
2. 두 원소의 합을 구한다. sum = aSum[left] + bSum[right]
   1. sum == t  t와 일치한다면 각 중복되는 원소의 크기를 구해 카운트 해준다. 
   2. sum > t   t보다 크다면 수가 작아져야 되므로 right --;
   3. sum < t   t보다 작다면 수가 커져야 하므로 left++;

 

이진탐색

투 포인터와 구하는 값은 같지만 탐색 방법만 다르다. 해당 방법은 upper_bound와 lower_bound를 사용하여 a, t-a의 두 원소의 값과 매칭되는 각 aSum구간과 bSum구간을 구한다. 

1. aSum을 기준으로 값을 설정한다. av = aSum[i];
2. aSum에서 av와 일치하는 구간을 구해준다. 
   1. aTerm = upper_bound(aSum, av) -lower_bound(aSum, av);
3. bSum에서 t-av와 일치하는 구간을 구해준다.
   1. bTerm = upper_bound(bSum, t-av)-  lower_bound(bSum, t-av);

 

 

투 포인터 풀이 코드 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {

	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		long t = Long.parseLong(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());
		}
		
		for(int i=1; i<n; i++) {
			a[i] += a[i-1];
		}
		for(int i=1; i<m; i++) {
			b[i] += b[i-1];
		}
		
		int aSize = n*(n+1)/2;
		int bSize = m*(m+1)/2;
		long[] aSum = new long[aSize];
		int idx=0;
		for(int i=0; i<n; i++) {
			for(int j=i; j<n; j++) {
				int av = a[j];
				if(i>0) av -= a[i-1];
				aSum[idx++] = av;
			}
		}
		
		long[] bSum = new long[bSize];
		idx=0;
		for(int i=0; i<m; i++) {
			for(int j=i; j<m; j++) {
				int bv = b[j];
				if(i>0) bv -= b[i-1];
				bSum[idx++] = bv;
			}
		}
		
		Arrays.sort(aSum);
		Arrays.sort(bSum);
		int left =0;
		int right = bSize-1;
		long cnt=0;
		while(left<aSize&& right>-1) {
			long asv = aSum[left], bsv = bSum[right];
			long sum = asv + bsv;
			if(sum ==t) {
				long ac =0, bc =0;
				while(left<aSize && asv == aSum[left]) {
					left++;
					ac++;
				}
			
				while(right>-1 && bsv == bSum[right]) {
					right--;
					bc++;
				}
				cnt += ac*bc;
			}
			if(sum>t) {
				right--;
			}else if(sum<t) {
				left++;
			}
		}
		System.out.println(cnt);
		
	}
}

 

이진탐색 풀이 코드 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {

	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		long t = Long.parseLong(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());
		}
		int aSize = n*(n+1)/2;
		int bSize = m*(m+1)/2;
		long[] aSum = new long[aSize];
		int idx=0;
		for(int i=0; i<n; i++) {
			int av=0;
			for(int j=i; j<n; j++) {
				av+=a[j];
				aSum[idx++] = av;
			}
		}
		
		long[] bSum = new long[bSize];
		idx=0;
		for(int i=0; i<m; i++) {
			int bv = 0;
			for(int j=i; j<m; j++) {
				bv += b[j];
				bSum[idx++] = bv;
			}
		}
		Arrays.sort(aSum);
		Arrays.sort(bSum);
		long cnt=0;
		for(int i=0; i<aSize;) {
			long av = aSum[i];
			long aTerm = upper_bound(aSum, av) -lower_bound(aSum, av);
			long bTerm = upper_bound(bSum, t-av)-  lower_bound(bSum, t-av);
			cnt+= aTerm*bTerm;
			i+=aTerm;
		}
		System.out.println(cnt);
	}
	
	static int upper_bound(long[] arr, long t) {
		int left = 0, right =arr.length; 
		while(left<right) {
			int mid = (left+right)/2;
			
			if(t >= arr[mid]) {
				left = mid+1;
			}else {
				right =mid;
			}
		}
		return right;
	}
	
	static int lower_bound(long[] arr, long t) {
		int left = 0, right =arr.length; 
		while(left<right) {
			int mid = (left+right)/2;
			
			if(t <= arr[mid]) {
				right =mid;
			}else {
				left = mid+1;
			}
		}
		return right;
	}
}

 

실행결과

투 포인터

이진탐색