Dev/PS

[백준] 1208 부분수열의 합 2 java

풋데브 2022. 4. 12. 13:08

import java.io.*;
import java.util.*;

public class Main {
    static BufferedReader br;
    static StringTokenizer st;
    static int n, s;
    static int[] arr;
    static List<Long> left;
    static List<Long> right;
    static long ans;

    public static void main(String[] args) throws IOException{
        input();
        dfs(left, arr.length/2, 0, 0);
        dfs(right, n, arr.length/2, 0);
        Collections.sort(left);
        Collections.sort(right);
        System.out.println(getCnt());
    }

    static void input() throws IOException {
        br = new BufferedReader(new InputStreamReader(System.in));
        st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        s = Integer.parseInt(st.nextToken());
        st = new StringTokenizer(br.readLine());
        arr = new int[n];
        left = new ArrayList<>();
        right = new ArrayList<>();
        for(int i = 0; i < n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }
        ans = 0;
    }

    static void dfs(List<Long> tmpArr, int n, int idx, long sum) {
        if(idx == n) {
            tmpArr.add(sum);
            return;
        }

        dfs(tmpArr, n, idx+1, sum);
        dfs(tmpArr, n, idx+1, sum+arr[idx]);
    }

    static long getCnt() {
        int start = 0;
        int end = right.size()-1;

        while(true) {
            if(start >= left.size() || end < 0) break;
            long leftVal = left.get(start);
            long rightVal = right.get(end);
            long sum = leftVal + rightVal;

            if(sum == s) {
                long leftCnt = 0;
                while(start < left.size() && left.get(start) == leftVal) {
                    leftCnt++;
                    start++;
                }

                long rightCnt = 0;
                while(end >= 0 && right.get(end) == rightVal) {
                    rightCnt++;
                    end--;
                }
                ans += leftCnt * rightCnt;
            }

            if(sum < s) start++;
            if(sum > s) end--;
        }

        return s == 0 ? ans-1 : ans;
    }
}

전에 풀었던 부분수열의 합과 동일한 문제이다. 다만 n의 범위가 전에는 20 이었지만 이 문제는 40으로 2배 늘었다.

그래서 재귀적인 방식으로 풀게되면 2^40 으로 시간초과가 나게된다. 그렇기 때문에 다른 방법을 써야한다.

 

나 역시 처음에 풀지 못하여 다른 분들의 블로그를 참고했다. 그리고 매번 내가 풀지 못하는 문제풀이를 볼 때 마다 이렇게도 풀 수 있구나 감탄하며 코드를 감상(?) 하게된다. 역시 다른 분들의 코드를 보는게 많이 도움이 된다.

 

이 문제의 풀이흐름은 아래와 같다.

 

1. 입력으로 받은 길이가 n인 배열을 받은 후 2개로 나눈다.

 

2. 각각 left, right 라고 하고 배열마다 재귀적으로 모든 경우의 수열을 저장한다. (2^20이니 시간안에 들어옴)

 

3. left, right를 오름차순 정렬해준다.

 

4. 투포인터 알고리즘으로 두 배열의 부분합 중 s와 같은 경우의 수를 모두 찾는다.

 

이 때, 4번의 경우에서 한 배열의 합 만으로 s가 되는 경우가 있는데 이 케이스는 각 배열마다 부분수열을 구할 때 공집합도 포함됬기 때문에 이 케이스 역시 구할 수가 있게된다.

 

그리고 s와 합이 같은 경우를 찾은 경우 left, right배열을 선형탐색하며 동일한 값이 있는지 개수를 샌 후 모두 카운트에 더해준다.

 

마지막으로 s가 0인 경우는 공집합 + 공집합의 경우가 포함되기 때문에 카운트에서 -1을 해주면 된다.