CodeForce - Number of Ways(466C)

문제(링크)

문제 설명:
N개의 숫자를 세 부분으로 나눴을 때 나눈 3개의 부분의 합이 모두 같을 경우의 수를
구하는 문제입니다.

풀이:
만약 세 구간이 위 조건을 만족한다면 다음과 같이 나타낼 수 있습니다.

1. 첫 번째 구간의 합 = A
2. 처음부터 두 번째 구간의 끝까지의 합 = 2A
3. 처음부터 세 번째 구간의 끝까지의 합 = 3A

위에서 3번의 경우는 N개의 수에 대해 prefix sum을 구했을 때 sum[N] 으로 나타내어
집니다.
즉, sum[N]이 3의 배수일 경우에만 조건을 만족할 수 있습니다.

sum[N]이 3의 배수일 경우부턴 다음과 같은 패턴을 찾습니다.
(A = sum[N] / 3 으로 구할 수 있습니다.)
맨 왼쪽부터 prefix sum을 탐색하면서 A , 2A , 3A 의 순서인 경우의 수를 세어주면 됩니다.

#include<iostream>
#include<string>
#include<algorithm>
#include<math.h>
#include<vector>
#include<queue>
#include<stdio.h>
using namespace std;
#define FIO ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL)
#define pii pair<int, int>

int N;
long long sum[500001];

int main(void) {
 FIO;

 cin >> N; 

 for (int i = 1; i <= N; i++)
 {
  long long v;
  cin >> v;
  sum[i] = sum[i - 1] + v;
 }

 if (sum[N] % 3 == 0) {
  long long A = sum[N] / 3;
  long long A_cnt = 0;
  long long cnt = 0;

  for (int i = 1; i < N; i++)
  {
   if (sum[i] == A * 2)
    cnt += A_cnt;
   if (sum[i] == A)
    A_cnt++;
  }
  cout << cnt;
 }
 else
  cout << 0;
}