백준 2143번 두 배열의 합 - 스위프트(Swift) 풀이

---

 

1. A의 모든 부 배열에 대해 합을 구해 각 합을 만들 수 있는 경우의 수를 카운트해둔다.
2. B의 모든 부 배열에 대해 합을 구해 각 합을 만들 수 있는 경우의 수를 카운트해둔다.
3. A의 모든 부 배열합 sum에 대해 (A의 부 배열 합이 sum이 되는 경우의 수) X (B의 부 배열 합이 T-sum이 되는 경우의 수)를 더한다.

 

1. A의 모든 부 배열에 대해 합을 구해 각 합을 만들 수 있는 경우의 수를 카운트해둔다.

 N=1000이므로 O(N^2)으로 모든 부 배열을 다 고려해볼 수 있다. 모든 부 배열에 대해 합을 구하면서 그 합을 만들 수 있는 경우의 수를 카운트해준다. 부 배열 합이 sum이 되는 경우의 수(count)를 딕셔너리에 dict[sum] = count 형태로 저장해준다. 원소 범위가 10^9이기 때문에 배열로 저장하는 것은 메모리 초과로 불가능하기 때문에 딕셔너리를 사용하였다.

 

2. B의 모든 부 배열에 대해 합을 구해 각 합을 만들 수 있는 경우의 수를 카운트해둔다.

1번 과정을 B 배열에 대해서도 반복하여 각 배열로 만들 수 있는 부 배열합을 모두 구해둔다.

 

3. A의 모든 부 배열합 sum에 대해 (A의 부 배열 합이 sum이 되는 경우의 수) X (B의 부 배열 합이 T-sum이 되는 경우의 수)를 더한다.

 A의 부 배열과 B의 부 배열의 모든 경우를 각각 구해두었으므로 이제 두 부 배열을 합치기만 하면 된다. A의 부 배열합이 sum인 경우에는 B의 부 배열합이 T-sum인 경우를 찾아 서로를 곱해주면 합이 T가 되는 경우의 수를 구할 수 있다. 이미 이것들을 딕셔너리에 dict[sum] = count 형태로 저장해두었다.

 전체 경우의 수는 A 딕셔너리의 모든 (키, 값) 쌍을 기준으로 탐색하면서 B의 부 배열합이 T-sum인 경우를 O(1)에 구해 곱한 뒤 더해주면 된다. 따라서 전체 시간복잡도는 O(N^2)이다.

반응형

---

import Foundation

var T = 0, n = 0, m = 0
var A = [Int]()
var B = [Int]()
var dictA = [Int:Int]()
var dictB = [Int:Int]()

func considerAllCases() {
    for start in 0..<n {
        var sum = 0
        
        for end in start..<n {
            sum += A[end]
            
            if dictA[sum] == nil {
                dictA[sum] = 1
            } else {
                dictA[sum]! += 1
            }
        }
    }
    
    for start in 0..<m {
        var sum = 0
        
        for end in start..<m {
            sum += B[end]
            
            if dictB[sum] == nil {
                dictB[sum] = 1
            } else {
                dictB[sum]! += 1
            }
        }
    }
}

func makeT() -> Int64 {
    var count: Int64 = 0
    
    for (key, val) in dictA {
        guard let cases = dictB[T - key] else { continue }
        
        count += Int64(val) * Int64(cases)
    }
    
    return count
}

func solution() {
    T = Int(readLine()!)!
    
    n = Int(readLine()!)!
    A = readLine()!.split(separator: " ").map{ Int(String($0))! }
    
    m = Int(readLine()!)!
    B = readLine()!.split(separator: " ").map{ Int(String($0))! }
    
    considerAllCases()
    print(makeT())
}

solution()