[백준] 2143 두 배열의 합 (C++)

백준 2143 두 배열의 합 문제
백준 2143 두 배열의 합 소스코드

📄 문제 설명

Problem

배열 A[1], A[2], …, A[n]에 대해서,
부 배열은 i부터 j까지의 배열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(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

Input

첫째 줄 : 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을 넘지 않는 정수이다.)

Output

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

Example Input

5
4
1 3 1 2
3
1 3 2

Example Output

7

📝 문제 해설

A라는 배열과 B라는 배열의 부배열을 더했을 때 T가 되는 case 개수를 구하는 문제이다.
이 문제를 완전 탐색으로 해결한다면 시간 복잡도가 O(N⁴)이 되므로 시간초과가 된다.

따라서 이 문제는 배열 A의 부배열합 <subA>와 배열 B의 부배열합 <subB>를 만든 후
두 배열을 가리키는 포인터를 만들어 문제를 해결하는 <Two Pointers>를 활용해 해결하였다.
이 방식으로 해결하면 시간 복잡도는 O(N²)이 된다.
(배열의 합을 구하는 경우 - O(N²), 배열 정렬 - O(Nlog N), 투포인터 - O(2N))

7453 합이 0이 되는 네 정수 (투포인터 설명)

</> Source Code

#include<bits/stdc++.h>

using namespace std;

bool comp(const int& a, const int& b){
	return a > b;
}

int main(){	
	int T, n, m;
	vector<int> A, B;
	vector<long> subA, subB;
	
	scanf("%d %d", &T, &n);	
	for(int i = 0; i < n; i++){
		int x;
		scanf("%d", &x);
		A.push_back(x);
	}
	
	scanf("%d", &m);
	for(int i = 0; i < m; i++){
		int x;
		scanf("%d", &x);
		B.push_back(x);
	}
	
	for(int i = 0; i < n; i++){
		long sum = 0;
		for(int j = i; j < n; j++){
			sum += A[j];
			subA.push_back(sum);
		}
	}
	
	for(int i = 0; i < m; i++){
		long sum = 0;
		for(int j = i; j < m; j++){
			sum += B[j];
			subB.push_back(sum);
		}
	}
	
	int ptA = 0, ptB = 0;
	long result = 0;
	
	sort(subA.begin(), subA.end());
	sort(subB.begin(), subB.end(), comp);
	
	while(ptA < subA.size() && ptB < subB.size()){
		long currentA = subA[ptA];
		long target = T - currentA;
		
		if(subB[ptB] > target){
			ptB ++;
		}
		else if(subB[ptB] == target){
			long countA = 0;
			long countB = 0;
			
			while(ptA < subA.size() && subA[ptA] == currentA){
				ptA ++;
				countA ++;
			}
			
			while(ptB <subB.size() && subB[ptB] == target){
				ptB ++;
				countB ++;
			}
           
			result += countA * countB;
		}
		else{
			ptA ++;
		}
	}
			
	printf("%ld", result);
	
	return 0;
}