codefores 466C

2019-02-15

466C Number of Ways 이 문제를 보고 떠오르는게 전체탐색 밖에 없기도 했고 Problem tags에도 brute force가 있어서 전탐으로 짜서 제출했더니 시간초과가 났다.

아무리 고민해도 가지치기 조금 하는 정도여서 다른 분의 해설을 보았다. Problem 466C 풀이
처음에는 저 풀이의 마지막 문단에 있는 Why i+2? Because we need a… 를 읽고도 바로 이해가 안갔다. 내 코드는 for문 도는 방향이 반대이므로 Why i-2? 가 되겠지만 이해하는데 큰 지장은 없을 것 이다.

먼저, if(sum % 3 == 0)블록 안에 있는 첫 번째 for문은 첫번째 뭉치인 0부터 i번째까지의 합에서, sum/3을 충족하는 경우가 몇 가지 있느냐? 이다.
두 번째 for문은 뒤에서 부터 돌면서 세번째 뭉치의 합이 sum/3인 경우를 찾는다. 아래는 example 01을 그림으로 나타낸 것이다. i-2인 이유는 위 링크에 있는 해설 그대로 무조건 3 뭉치를 만들기 때문이다.

그러니까 i도 안되고, i-1도 안된다. 만약 문제에서 연속된 세 부분합이 아닌 두 부분합을 구하는 방법의 수라면 cnt[i-1]일 것이다. cnt[i-2]는 인덱스 0..i-2까지 왔을 때, sum/3을 만족하는 첫 번째 뭉치의 최대 갯수이고, i-1은 두 번째 뭉치, i는 i..n-1까지의 합인 세 번째 뭉치이다. 이 것을 그림으로 그려보면 아래와 같다. 빨간색으로 묶어놓은 부분이 실제로 분할되는 영역이다.

for문을 천천히 따라 가 보자. a[]를 뒤에서 부터 앞으로 돌다가 sum/3을 만족시키는 곳을 만났다. 그 지점을 회색으로 표시했다. i-2가 실제로 a[]를 나누는 첫 번째 인덱스가 아니어도 상관이 없다. 왜냐하면 우리는 3번째 뭉치가 sum/3일 때, 만들 수 있는 첫 번째 뭉치가 최대 몇 개 인지를 원하기 때문이다. 이걸 알면 두번째 뭉치는 고정이다.

466c example 1

계속 for문이 돈다. sum/3을 만족시키는 곳을 또 만났다. 그 지점을 회색으로 표시했다.

466c example 1

ans = 1 + 1이므로 example 01의 output이 2가 되는 것이다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        int n = Integer.parseInt(br.readLine());
        int[] a = new int[n];
        long sum = 0;

        StringTokenizer st = new StringTokenizer(br.readLine());
        for(int i = 0; i < n; i++) {
            a[i] = Integer.parseInt(st.nextToken());
            sum += a[i];
        }

        if(sum % 3 == 0) {
            long partSum = sum / 3;
            int cnt[] = new int[n];

            long buff = 0;
            // 첫 번째 for 문
            for(int i = 0; i < n; i++) {
                if(i >= 1) cnt[i] = cnt[i - 1];

                buff += a[i];
                if(buff == partSum) cnt[i]++;
            }

            buff = 0;
            long ans = 0;
            // 두 번째 for 문
            for(int i = n - 1; i >= 2; i--) {
                buff += a[i];
                if(buff == partSum) ans += cnt[i - 2];
            }

            System.out.println(ans);
        }
        else {
            System.out.println(0);
        }
    }
}