[6개월 안에 백준 플래티넘 달성하기] 23. 두 배열의 합 2143

두 배열의 합

2143번: 두 배열의 합

 

요즘 들어 문제를 못 푸는 경우가 많아졌습니다. 물론 정말 어려워서 해설을 보고도 이해하는데 시간이 걸리는 문제도 있지만 대부분의 문제는 해설을 보는 순간 왜 이 간단한 생각을 못했지라는 후회와 현타가 오는 문제입니다.

두 배열의 합 문제도 이러한 문제 중에 하나입니다.

 

문제

 

풀이

역시나 가장 먼저 DP로 풀이를 생각해 봤습니다. 2차원 배열을 이용한 DP인가 싶어서 자연스럽게 2차원 배열에 값을 끼어넣어주는 과정에서 가장 간단한 모든 경우를 구하는 방법을 생각했습니다.

 

A배열에서 부 배열의 합의 개수를 구하는 경우의 수는 (n*(n+1))  / 2로 n=1000일 때 최대 50만의 시간 복잡도를 가집니다.

B배열도 마찬가지입니다.

(N*(N+1))  / 2   는   1~n까지의 합 공식이다.

 

이제 A의 부 배열의 합과 B의 부 배열의 합을 더했을 때 T가 되는 경우를 구하면 됩니다.

가장 간단한 방법은 이중 for문을 돌면서 모든 경우를 체크해 주는 것입니다.

for (i in A의 부분배열의 합)
    for (j in B의 부분배열의 합) 
        if (i + j == T) result++

 

하지만 위와 같이 짤 경우 시간 복잡도는 최대 50만*50만으로 시간초과가 발생합니다.

 

여기서 저는 해결법을 생각하지 못하였습니다. 지금 생각해 보면 정말 왜 이 방법을 생각하지 못하였는지 모르겠습니다.

해결법은 이분 탐색, 혹은 투 포인터입니다.

 

A의 부분배열의 합에 더해서 T가 되는 B의 부분배열의 합을 이분탐색 혹은 투 포인터로 찾으면 시간초과가 발생하지 않습니다.

• 이분 탐색의 경우 (50만 * log 2의 50만)
• 투 포인터의 경우 (50만+50만)

(추가적으로 글 마지막에 map을 이용하여 더 간단하게 풀 수 있는 코드가 있습니다.)

수의 범위가 크기 때문에 Long타입으로 지정해 줘야 됩니다.
결과값 또한 Long타입으로 지정해 줘야 됩니다.

---

먼저 이분 탐색 풀이입니다.

 

이분 탐색으로 풀 경우 한 가지 주의점이 있습니다.

이분탐색 시에 T가 되는 B의 부분 배열의 합이 여러 개일 경우에도 하나의 답만 내놓기 때문에 이를 처리해 줘야 됩니다.

여기서 저는 map을 이용하여 해결해 주었습니다. key값에는 기존 원소를 넣고 value에는 원소의 개수를 넣어주었습니다.

 

import java.io.StreamTokenizer

fun main() = StreamTokenizer(System.`in`.bufferedReader()).run {
    fun input(): Long {
        nextToken()
        return nval.toLong()
    }
    val T = input()
    val aArr = LongArray(input().toInt()) {input()}
    val bArr = LongArray(input().toInt()) {input()}

    val aSum = arrayListOf<Long>()
    val bTreeMap = hashMapOf<Long, Long>() // 합, 개수

    // A의 부 배열의 합을 모두 구한다.
    for (i in aArr.indices) {
        var sum = 0L
        for (j in i until aArr.size) {
            sum += aArr[j]
            aSum.add(sum)
        }
    }
    
    // B의 부 배열의 합을 모두 구한다.
    // A와 다르게 개수 중복을 허용하지 않는다.
    for (i in bArr.indices) {
        var sum = 0L
        for (j in i until bArr.size) {
            sum += bArr[j]
            bTreeMap[sum] = bTreeMap.getOrDefault(sum, 0) + 1
        }
    }

    // 이진 탐색
    val keys = bTreeMap.keys.toLongArray().sorted() // 이진 탐색을 위해 정렬
    fun binarySearch(num: Long): Long {

        var left = 0
        var right = keys.size-1

        while (left <= right) {
            var mid = (right+left)/2

            if (keys[mid] == num) {
                return bTreeMap[num]!! // 해당 원소 개수를 반환
            } else if (keys[mid] < num) {
                left = mid + 1
            } else {
                right = mid - 1
            }
        }
        return 0
    }

    // A의 부분 배열의 합 + B의 부분 배열의 합 = T 가 되는 경우를 이진탐색으로 찾는다.
    var result = 0L
    for (num in aSum) {
        val b = binarySearch(T-num)
        result += b
    }
    print(result)
}

 

---

투 포인터 풀이입니다.

 

따로 중복 원소를 처리하기 위해 map을 사용하지 않아도 됩니다. 

A와 B의 부분 배열의 합을 담은 배열을 정렬한 뒤 A의 왼쪽 인덱스, B의 오른쪽 인덱스로부터 원소들의 합을 T와 비교하여 포인터를 이동시켜 주면 됩니다.

 

import java.io.StreamTokenizer

fun main() = StreamTokenizer(System.`in`.bufferedReader()).run {
    fun input(): Long {
        nextToken()
        return nval.toLong()
    }
    val T = input()
    val aArr = LongArray(input().toInt()) {input()}
    val bArr = LongArray(input().toInt()) {input()}

    val aSum = arrayListOf<Long>()
    val bSum = arrayListOf<Long>()

    // A의 부 배열의 합을 모두 구한다.
    for (i in aArr.indices) {
        var sum = 0L
        for (j in i until aArr.size) {
            sum += aArr[j]
            aSum.add(sum)
        }
    }

    // B의 부 배열의 합을 모두 구한다.
    for (i in bArr.indices) {
        var sum = 0L
        for (j in i until bArr.size) {
            sum += bArr[j]
            bSum.add(sum)
        }
    }

    // 투 포인터
    var result = 0L
    var left = 0            // A
    var right = bSum.size-1 // B

    aSum.sort()
    bSum.sort()

    while (left < aSum.size && right >= 0) {
        val sum = aSum[left] + bSum[right]

        // T와 같을 경우
        if (sum == T) {
            var at = 0L
            var bt = 0L
            val tLeft = aSum[left]
            while (left < aSum.size && tLeft == aSum[left]) { // 중복 숫자 처리
                left++
                at++
            }
            val tRight = bSum[right]
            while (right > -1 && tRight == bSum[right]) { // 중복 숫자 처리
                right--
                bt++
            }
            result += at * bt
        } else if (sum < T) {
            left++
        } else if (sum > T) {
            right--
        }
    }

    print(result)
}

 

---

근데 사실 위 두 방식보다 더 쉽고 빠른 풀이가 있습니다.

A의 부분배열의 합을 map에 넣고 B의 부분배열의 합을 구할 때 바로바로 비교하는 방법입니다.

fun main() = StreamTokenizer(System.`in`.bufferedReader()).run {
    fun input(): Int {
        nextToken()
        return nval.toInt()
    }
    val T = input()
    val aArr = LongArray(input()) {input().toLong()}
    val bArr = LongArray(input()) {input().toLong()}

    val aTreeMap = hashMapOf<Long, Long>()

    for (i in aArr.indices) {
        var sum = 0L
        for (j in i until aArr.size) {
            sum += aArr[j]
            aTreeMap[sum] = aTreeMap.getOrDefault(sum, 0) + 1
        }
    }

    var result = 0L
    for (i in bArr.indices) {
        var sum = 0L
        for (j in i until bArr.size) {
            sum += bArr[j]
            result += aTreeMap.getOrDefault(T-sum, 0) // map을 통해 바로 찾을 수 있다.
        }
    }
    print(result)
}

 

정말 간단한 풀이인데 생각해 내기 쉽지가 않습니다.