두 배열의 합

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

각 배열의 모든 부분합들을 미리 구하여 저장해놓는다.

이렇게 구해진 두 개의 부분합 배열을 각각 CAND_A, CAND_B라고 하자.

두 부분합 배열을 정렬한 뒤에 서로 다른 방향에서 시작하는 투 포인터 아이디어를 적용해주면 답을 쉽게 구할 수 있다.

단, 주의할 점으로는 투 포인터 아이디어는 정렬한 배열이 같은 원소를 포함하게 되는 경우 제대로 카운팅을 하기 어려워지는 문제가 있다는 점이다.

예를 들어, CAND_A, CAND_B가 아래와 같이 주어지고 lo, hi가 아래와 같다고 해보자.

• CAND_A: ... 3(lo) 3 3 ...
• CAND_B: ... 3 3 3(hi) ...

이 경우 T=6이라면, lo 또는 hi를 이동시키는 아이디어로는 올바르게 경우의 수를 셀 수 없다.

따라서, 나는 중복을 허용하지 않도록 CAND_A, CAND_B를 값, 빈도의 map으로 관리하였다.

#include <cstdio>
#include <cstdint>
#include <vector>
#include <algorithm>
#include <map>

using namespace std;

int T;
int N;
int M;
vector<int> PSUM_A;
vector<int> PSUM_B;
map<int, int> CAND_A;
map<int, int> CAND_B;

int main(void)
{
  scanf(" %d", &T);

  scanf(" %d", &N);
  PSUM_A.emplace_back(0);
  for (int it = 1; it <= N; ++it)
  {
    int a;
    scanf(" %d", &a);
    PSUM_A.emplace_back(a + PSUM_A[it - 1]);
  }

  scanf(" %d", &M);
  PSUM_B.emplace_back(0);
  for (int it = 1; it <= M; ++it)
  {
    int b;
    scanf(" %d", &b);
    PSUM_B.emplace_back(b + PSUM_B[it - 1]);
  }

  for (int i = 1; i <= N; ++i)
  {
    for (int j = i; j <= N; ++j)
    {
      ++CAND_A[PSUM_A[j] - PSUM_A[i - 1]];
    }
  }

  for (int i = 1; i <= M; ++i)
  {
    for (int j = i; j <= M; ++j)
    {
      ++CAND_B[PSUM_B[j] - PSUM_B[i - 1]];
    }
  }

  int64_t answer = 0;
  auto lo = CAND_A.begin();
  auto hi = CAND_B.rbegin();

  while (lo != CAND_A.end() && hi != CAND_B.rend())
  {
    int sum = lo->first + hi->first;
    if (sum == T)
    {
      answer += (int64_t)lo->second * (int64_t)hi->second;
      ++lo;
      ++hi;
    }
    else if (sum < T)
    {
      ++lo;
    }
    else
    {
      ++hi;
    }
  }

  printf("%ld\n", answer);

  return 0;
}