백준 - 두 배열의 합 (2143)

2143번: 두 배열의 합

문제

배열 A의 부분 구간의 합과 배열 B의 부분 구간의 합, 이 두개의 합이 T와 일치하는 경우를 구하는 문제 입니다.

접근

각 배열 A,B에서의 부분합이 될 수 있는 수들을 나열합니다. 부분합이 될 수 있는 수들을 모은 배열은 각각 SA, SB라고 합시다. 그렇다면 SA를 순회하면서 SB의 요소 중에 T-SA[i]와 일치하는 값을 찾으면 됩니다.

ans <- 1
for i from 0 to (Length of SA)-1 do
begin
	for j from 0 to (Length of SB)-1 do
    begin
    	if T-SA[i] == SB[j] then ans <- (ans + 1)
    end
endFor
print ans

그러나 이는 for문이 두번 들어가기 때문에 O(SA*SB)시간이 걸립니다. SA, SB의 최대 길이는 500000이기 때문에 TLE가 발생합니다. 하지만 SB를 정렬한 다음 SB를 이분탐색으로 일치하는 값을 구한다면 O(SA*log2(SB))이 걸리기 때문에 속도가 훨씬 빨라집니다. 단, 찾고자 하는 값이 여러개가 있으므로 upper_bound와 lower_bound를 사용햐여 찾고자 하는 값의 범위를 구합니다.

ans <- 0
sort(SB)

for i from 0 to (length of SA) - 1 do
begin
	target <- T - SA[i]
    l = lower_bound(SB, target)
    r = upper_bound(SA, target)
    ans <- ans + (r - l)
endFor

print(ans)

수도 코드

A <- 입력된 배열 A
B <- 입력된 배열 B
T <- target
ans <- 0

# 구간합 구하기
CA <- A의 구간합
CB <- B의 구간합

# 부분합의 경우의 수 구하기 (CA, CB를 이용하여 계산)
SA <- A의 부분합 리스트
SB <- B의 부분합 리스트

sort SB
for i from 0 to length of SA - 1 do
begin
	target <- T - SA[i]
    l <- lower_bound(SB, target)
    r <- lower_bound(SB, target)
    ans <- ans + (r - l)
end

print(ans)

 

코드(C)

#include <stdio.h>
#include <stdlib.h>
#define MAX 1000
#define CASE_MAX 1000000

int compare(const void* o1, const void* o2)
{
    int n1 = *(int*)o1;
    int n2 = *(int*)o2;
    
    if(n1 > n2) return 1;
    else if(n1 < n2) return -1;
    else return 0;
}

int main(void)
{
	int t, n, m, v, l, r;
	register int i,j,k,ak,bk;
	int a[MAX], b[MAX];
	int sa[CASE_MAX], sb[CASE_MAX];
	long long ans = 0;

	scanf("%d", &t);
	scanf("%d", &n);
	for(i=0;i<n;i++) 
	{ 
		scanf("%d", &v); 
		a[i] = (i==0 ? v : a[i-1] + v);
	}
	
	scanf("%d", &m);
	for(i=0;i<m;i++) 
	{ 
		scanf("%d", &v); 
		b[i] = (i==0 ? v : b[i-1] + v); 
	}

	for(i=0;i<CASE_MAX;i++) sa[i] = sb[i] = 1e9+7;

	k=0;
	for(i=0;i<n;i++)
		for(j=i;j<n;j++)
			sa[k++] = a[j] - (i==0 ? 0 : a[i-1]);
	ak=k;
	k=0;
	for(i=0;i<m;i++)
		for(j=i;j<m;j++)
			sb[k++] = b[j] - (i==0 ? 0 : b[i-1]);
	bk=k;

	// sort sb
	qsort(sb, bk, sizeof(int), compare);
	for(k=0;k<ak;k++)
	{
		int v = t - sa[k];	// 찾아야 할 값
		l = 0, r = bk;
		long long lc, rc;

		// lower bound
		while(l<r)
		{
			int m = (l+r)>>1;
			if(sb[m] < v) l = m+1;
			else r = m;
		}
		lc = r;

		// upper bound
		l = 0; r = bk;
		while(l<r)
		{
			int m = (l+r)>>1;
			if(sb[m] <= v) l = m+1;
			else r = m;
		}
		rc = r;
		ans += (rc - lc);
	}
	printf("%lld\n", ans);
	return 0;
}

lower bound와 upper bound를 직접 구현을 해봤는데, 잘못 구현하는 바람에 한참 애를 먹었습니다. lower bound와 upper bound, 그리고 이분탐색에 대해 연구를 해서 따로 포스팅을 할까 생각중입니다.

물론 C++이나 Python같은 언어에서는 직접 라이브러리로 지원을 해 주니, 빠른 시간에 풀고 싶으신 분들은 라이브러리를 직접 사용하시면 되겠습니다.