[백준(Baekjoon)][자바(java)] 2143 : 두 배열의 합 / 이분 탐색

 

www.acmicpc.net/problem/2143

2143번: 두 배열의 합

 

문제 풀이

 

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

public class Main {
	
	public static int lower_bound( int val, int start, int end, int arr[] ) {	
		// 배열 a에서 k값이 처음으로 나타나는 인덱스
		while( start < end ) {
			int mid = (start+end)/2;
			if( arr[mid] >= val )		
				end = mid;
			else start = mid+1;
		}
		return end;
	}
	
	public static int upper_bound( int val, int start, int end, int arr[] ) {	
		// 배열 a에서 k보다 큰 값이 처음으로 나타나는 인덱스
		while( start < end ) {
			int mid = (start+end)/2;
			if( arr[mid] > val )	
				end = mid;
			else start = mid+1;
		}
		return end;
	}

	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], i, j;
		StringTokenizer st = new StringTokenizer( br.readLine() );
		for( 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( i = 0; i < m; i++ )
			B[i] = Integer.parseInt( st.nextToken() );
		br.close();
		
		int len_A = 0, len_B = 0;
		for( i = 1; i <= n; i++ ) len_A += i;
		for( i = 1; i <= m; i++ ) len_B += i;
		int A_ss[] = new int[len_A]; // sub sum
		int B_ss[] = new int[len_B];	
		int idx = 0, ss;
		for( i = 0; i < n; i++ ) {
			ss = 0;
			for( j = i; j < n; j++ ) {
				ss += A[j];
				A_ss[idx++] += ss;
			}
		}
		idx = 0;
		for( i = 0; i < m; i++ ) {
			ss = 0;
			for( j = i; j < m; j++ ) {
				ss += B[j];
				B_ss[idx++] += ss;
			}
		}
		
		Arrays.sort( B_ss );
		
		long cnt = 0; int val, ub, lb;
		for( i = 0; i < len_A; i++ ) {
			val =  T - A_ss[i];
			ub = upper_bound( val, 0, len_B, B_ss );
			lb = lower_bound( val, 0, len_B, B_ss );
			cnt += ub - lb;
		}
		
		System.out.println( cnt );
	}
}

 

*  이분 탐색, 누적 합, 정렬

-  배열 A[], B[] 각각의 부분 합을 구해 A_ss[], B_ss[] 생성 ( sub sum )
-  B_ss[] 정렬 후 이분탐색 하면서 T - A_ss[i] 를 찾음
-  upper bound, lower bound를 구해 차이를 cnt에 더함