[JAVA] 백준 2143번 【G3.두 배열의 합】

문제

 

2143번: 두 배열의 합

---

풀이

각 배열의 부분합을 전부 구한다.

 

	    List<Long> sumsA = new ArrayList<>();
		n = Integer.parseInt(br.readLine());
        A = Arrays.stream(br.readLine().split(" ")).mapToLong(Long::valueOf).toArray();

        for (int i = 0; i < n; i++) {
            long res = 0;
            for (int j = i; j < n; j++) {
                res += A[j];
                sumsA.add(res);
            }
        }

 

각 배열의 최대 원소 개수는 1000개이므로 이중 포문을 사용해도 시간복잡도에 크게 문제가 없음을 알 수 있다.

 

그렇게, 배열 A와 배열 B의 부분합을 다 구한 배열을 정렬한 뒤 투 포인터 방식을 사용하면 된다.

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;


public class Main {
    int T;
    int n;
    int m;
    long[] A;
    long[] B;
    List<Long> sumsA = new ArrayList<>();
    List<Long> sumsB = new ArrayList<>();

    public static void main(String[] args) throws IOException {
        new Main().sol();
    }

    private void sol() throws IOException {
        setUp();
        Collections.sort(sumsA);
        Collections.sort(sumsB);
        
        long cnt = 0;
        int left = 0, right = sumsB.size() - 1;
        while (left <= sumsA.size()-1 && right >= 0) {
            long sum = sumsA.get(left) + sumsB.get(right);

            if (sum < T) {
                left++;
            } else if (sum > T) {
                right--;
            } else {
                long cntOfL = 1;
                long cntOfR = 1;

                while (left + 1 <= sumsA.size() - 1 && sumsA.get(left).equals(sumsA.get(left + 1))) {
                    cntOfL++;
                    left++;
                }

                while (right - 1 >= 0 && sumsB.get(right).equals(sumsB.get(right - 1))) {
                    cntOfR++;
                    right--;
                }

                cnt += cntOfL * cntOfR;

                left++;
                right--;
            }
        }
        System.out.println(cnt);

    }

    private void setUp() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        T = Integer.parseInt(br.readLine());
        n = Integer.parseInt(br.readLine());
        A = Arrays.stream(br.readLine().split(" ")).mapToLong(Long::valueOf).toArray();
        m = Integer.parseInt(br.readLine());
        B = Arrays.stream(br.readLine().split(" ")).mapToLong(Long::valueOf).toArray();

        for (int i = 0; i < n; i++) {
            long res = 0;
            for (int j = i; j < n; j++) {
                res += A[j];
                sumsA.add(res);
            }
        }

        for (int i = 0; i < m; i++) {
            long res = 0;
            for (int j = i; j < m; j++) {
                res += B[j];
                sumsB.add(res);
            }
        }
    }
}

 

 

이 문제를 풀면서 주의해야 할 점이 있었다.

 

				while (left + 1 <= sumsA.size() - 1 && sumsA.get(left) == sumsA.get(left + 1)) {
                    cntOfL++;
                    left++;
                }

                while (right - 1 >= 0 && sumsB.get(right) == sumsB.get(right - 1)) {
                    cntOfR++;
                    right--;
                }

 

처음에는 위 코드 처럼 whlie 조건문에 equals() 가 아닌 == 비교를 했었고

 

통과가 되지 않았다.

 

List<> 는 Wrapper 타입으로 보관하기 때문에 당연히 == 비교가 성립하지 않는다는 점을 간과했다.

 

예제로 나온 테스트 케이스는 통과해서 눈치를 못챘던 것이다.

 

chatGPT에 따르면 작은 값의 경우 Wrapper 타입이라도 값을 공유해서 == 비교도 true를 반환하고 값이 커지면 false를 반환한다고 한다.

 

따라서, 컬렉션에서 get() 해서 변수로 할당하고 비교를 하는 것이 아니면

 

컬렉션에서의 원소 비교는 꼭 equals()를 쓰도록 주의해야함을 알 수 있던 문제였다.