[BOJ 2143] 두 배열의 합

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

2143번: 두 배열의 합

---

n과 m의 값이 크지 않기 때문에 각 배열의 합을 구하는 것은 가능하다.

 

이를 이용하여 A배열과 B배열의 연속합을 저장해둔 배열을 만들고, 이를 정렬한 뒤 각 합을 따로 합쳐보면서 개수를 세면 된다.

 

sA = A배열에서 i ~ j ( 0<= i <= j <= n)까지의 합을 저장해둔 배열(정렬됨)

sB = B배열에서 i ~ j ( 0<= i <= j <= m)까지의 합을 저장해둔 배열(정렬됨)

 

합을 저장해둔 배열의 양 끝을 가리키는 변수를 만들고 이를 이동시키면서 합이 T가 되는 경우의 수를 계산하면 된다.

 

sA은 처음부터 확인하고, sB는 끝부터 확인을 한다면 값이 작다면 sA를 가리키는 값을 증가시키면 되고, 크다면 sB를 가리키는 값을 감소시키면서 확인하면 된다.

 

sA와 sB는 정렬되어 있기 때문에 가능하다.

---

#include <stdio.h>
#include <vector>
#include <algorithm>

using namespace std;

int T;
int n, m;
int a[1001];
int b[1001];
unsigned long long int ans;

vector<int> sA, sB;

int main()
{
	scanf("%d", &T);
	scanf("%d", &n);
	for (int i = 0; i < n; i++)
		scanf("%d", a + i);

	scanf("%d", &m);
	for (int j = 0; j < m; j++)
		scanf("%d", b + j);

	for (int i = 0; i < n; i++)
	{
		int sum = 0;
		for (int j = i; j < n; j++)
		{
			sum += a[j];
			sA.push_back(sum);
		}
	}
	for (int i = 0; i < m; i++)
	{
		int sum = 0;
		for (int j = i; j < m; j++)
		{
			sum += b[j];
			sB.push_back(sum);
		}
	}
	sort(sA.begin(), sA.end());
	sort(sB.begin(), sB.end());

	int l = 0, r = sB.size() - 1;

	while (1)
	{
		if (sA[l] + sB[r] < T)
			l++;
		else if (sA[l] + sB[r] > T)
			r--;
		else if (sA[l] + sB[r] == T)
		{
			long long cnt1 = 1, cnt2 = 1;

			//sA[현재] == sA[현재 + 1]이면 l을 더해도 항상 T이기 때문에 
			while (l + 1 < sA.size() && sA[l] == sA[l + 1])
			{
				l++;
				cnt1++;
			}
			while (r - 1 >= 0 && sB[r] == sB[r - 1])
			{
				r--;
				cnt2++;
			}
			ans += cnt1 * cnt2;
			l++, r--;
		}


		if (l >= sA.size() || r < 0)
			break;
	}
	printf("%llu", ans);
}