쌉쌀한 무화과의 개발로그

[문제]

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

​

[예제]

>문제 풀이

A[ ]: 1 3 1 2

listA: 1, 4, 5, 7, 3, 4, 6, 1, 3, 2

sort: 1, 1, 2, 3, 3, 4, 4, 5, 6, 7

​

B[ ]: 1 3 2

listB: 1, 4, 6, 3, 5, 2

sort: 1, 2, 3, 4, 5, 6

​

List에 각 배열로 만들 수 있는 부분합을 저장해놓고, lower bound와 upper bound를 구하여

listA.get(i) + listB.get(?)==T인 개수를 구한다.

​

단, lower bound와 upper bound를 구하기 전에 sort를 해줘야 한다.

lower bound는 listA.get(i)+listB.get(lb)가 T이상인 값을 구하고

upper bound는 listA.get(i)+listB.get(ub)가 T초과인 값을 구한다.

그리고 cnt+= ub-lb;

+) 투 포인터 탐색 방법을 사용할 수도 있다.

+) map을 활용할 수도 있는데 그러면 코드가 더 간결해진다.

>전체 코드

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

public class Main {
	
	public static void main(String[] args) throws IOException {
		BufferedReader br= new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;
		
		int T, N, M; //-십억<=T<=십억 // N: 1000이하의 정수, M: 1000 이하의 정수
		long[] A, B; //각 원소는 절댓값 백만 이하
		List<Long> listA= new ArrayList<>();
		List<Long> listB= new ArrayList<>();
		
		//입력
		T= Integer.parseInt(br.readLine());
		N= Integer.parseInt(br.readLine());
		A= new long[N];
		st= new StringTokenizer(br.readLine());
		for(int i=0; i<N; i++) {
			A[i]= Long.parseLong(st.nextToken());
		}
		
		M= Integer.parseInt(br.readLine());
		B= new long[M];
		st= new StringTokenizer(br.readLine());
		for(int i=0; i<M; i++) {
			B[i]= Long.parseLong(st.nextToken());
		}
		
		//나올 수 있는 합 구해서 list에 넣기
		long sum=0;
		for(int i=0; i<N; i++) {
			sum= A[i];
			listA.add(sum);
			for(int j=i+1; j<N; j++) {
				sum+=A[j];
				listA.add(sum);
			}
		}
		listA.sort(null);
		
		for(int i=0; i<M; i++) {
			sum= B[i];
			listB.add(sum);
			for(int j=i+1; j<M; j++) {
				sum+=B[j];
				listB.add(sum);
			}
		}
		listB.sort(null);
		
		//계산 및 출력
		long cnt=0;
		int left, right, mid, lb, ub;
		for(int i=0; i<listA.size(); i++) {
			//lower
			left=0;
			right=listB.size();
			while(left<right) {
				mid=(left+right)/2;
				if(listA.get(i)+listB.get(mid)>=T){
					right= mid;
				}else {
					left= mid+1;
				}
			}
			lb= left;
			
			//upper
			left=0;
			right=listB.size();
			while(left<right) {
				mid=(left+right)/2;
				if(listA.get(i)+listB.get(mid)<=T){
					left= mid+1;
				}else {
					right= mid;
				}
			}
			ub= left;
			cnt+=ub-lb;
		}//for
		
		System.out.println(cnt);
	}//main
}

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