[Codeforces 1118B] Tanya and Candies (Prefix Sum,부분합)

[Codeforces 1118B] Tanya and Candies

n일동안 하루마다 할당된 캔디의 수가 주어진다. 이때 어떤 하루의 캔디를 모두 삭제 한 뒤의 상태(그 이후의 날은 앞으로 한 칸씩 당겨짐)에서 홀수 날 캔디의 합과 짝수 날 캔디의 합이 같은 날이 n일 중에 몇 일인지 찾는 문제이다.

naive하게 푼다면 i번째 날을 계산하기 위해서 1~(i-1)번 째까지 수 + (i+1)~n번 째까지의 캔디를 모두 더하면 되고, 이 풀이는 O(N^2)이다. 하지만 이 문제는 n이 최대 20만이기 때문에 시간이 초과될 것이다. 그래서 관건은 1~(i-1)번째까지, (i+1)~n번 째까지의 캔디의 합을 빨리 구해야한다. 당연하게도 Prefix Sum(부분합)으로 푸는 문제이다. 이 문제에서 한 가지 처리해주어야 할 것은 i번째 날의 캔디를 삭제했다면 그 이후의 날은 한 칸씩 당겨져서 홀짝이 뒤바뀌게 된다. 

예를 들어 예시 입력 1에서 7 / 5 5 4 5 5 5 6 이 있을 때 3번째 날(4)를 삭제한다면 그 결과는 5 5 5 5 5 6과 같아진다. 따라서 i번 째 날을 삭제한 뒤, 짝수 날의 합을 구할 때는 i번 째 날 전까지 짝수 날의 합 + i번 째 날 이후의 홀수 날의 합이 되야 한다. 같은 원리로 홀수 날의 합을 구할 때는 i번째 날 전까지 홀수 날의 합 + i번 째 날 이후의 짝수 날의 합이다. 이를 위해서 앞에서부터 누적해나가는 짝수 날의 합(ePre), 홀수 날의 합(oPre), 뒤에서부터 누적해나가는 짝수 날의 합(eSuf), 홀수 날의 합(oSuf)변수가 필요하고 매 루프가 끝날 때마다 갱신해주면 된다. 

#include<stdio.h>
#define SIZE 200009
int arr[SIZE];
int ePre,oPre,eSuf,oSuf;
int main(){
    int n;  scanf("%d",&n);
 
    for(int i=1 ; i<=n ; i++){
        int inp;    scanf("%d",&inp);
        arr[i]=inp;
        if(i&1){
            oSuf+=inp;
        }
        else eSuf+=inp;
    }
 
    int ans=0;
    for(int i=1 ; i<=n ; i++){
        int even=0,odd=0;
 
        if(i&1){
            oSuf-=arr[i];
        }
        else eSuf-=arr[i];
 
        even=ePre+oSuf,odd=oPre+eSuf;
 
        if(even==odd) ans++;
 
        if(i&1){
            oPre+=arr[i];
        }
        else ePre+=arr[i];
    } 
    printf("%d\n",ans);
    return 0;
}