2167번: 2차원 배열의 합

이중 반복문으로 풀기

문제 풀이 아이디어

단순하게 모든 case의 입력을 받아 바로바로 이중 반복문을 통해 더해서 문제를 푸는 방식입니다. 아주 간단한 방법이지만 🚫 시간 초과를 받을 가능성이 큽니다.

코드

// 입력 받기
let input = readLine()!.split(separator: " ").map { Int(String($0))! }
let N = input[0], M = input[1]

// 2차원 배열에 입력 저장하기
var matrix = [[Int]]()
for _ in 0..<N {
    matrix.append(readLine()!.split(separator: " ").map { Int(String($0))! })
}

// 각 case 별로 이중 반복문으로 합 구하기
let T = Int(readLine()!)!
for _ in 0..<T {
    let line = readLine()!.split(separator: " ").map { Int(String($0))! }
    let i = line[0], j = line[1], x = line[2], y = line[3]
    var sum = 0
    for r in i...x {
        for c in j...y {
            sum += matrix[r - 1][c - 1]
        }
    }
    print(sum)
}

DP로 풀기

문제 풀이 아이디어

이 문제의 특징 중에 하나는 반복해서 이차원 배열의 합을 구하고 있다는 것입니다. 한번만 구하는 것이라면 위 처럼 이중 반복문으로 구현하는 것도 괜찮습니다만 구해야 합은 K개로 최대 10000개의 입력이 주어집니다. 따라서 DP를 활용해서 미리 구한 합을 메모리에 저장해놓고 사용해봅시다.

DP를 사용하기 위해서는 2단계를 거쳐야 합니다. 단순히 (0, 0) ~ (r, c)의 합을 구하는 것이 아니라 (i, j) ~ (x, y)의 합을 구하는 것이므로 4차원 배열을 사용해야 한다고 생각할 수도 있지만 (i, j) ~ (x, y)의 합을 (0, 0) ~ (r, c)가 저장된 cache를 활용해서 구할 수 있다는 점을 파악하면 2차원 배열의 cache로도 충분히 해결할 수 있습니다.

(0, 0) ~ (r, c)의 합 구하기

dp(r, c) (= (0, 0) ~ (r, c)의 합)를 구하는 점화식을 그림으로 나타내면 다음과 같습니다. 먼저 각각 파란색과 초록색 빗금으로 표시된 dp(r - 1, c)와 dp(r, c - 1)을 더합니다. 그리고 나서 두번 더해진, 두 빗금이 겹치는 부분인 dp(r - 1, c - 1)은 빼줍니다. 그리고 아직 추가되지 않은 빨간 빗금인 matrix[r][c]를 더해줍니다.

즉 dp(r, c) = dp(r - 1, c) + dp(r, c - 1) - dp(r - 1, c - 1) + matrix[r][c]의 점화식을 구할 수 있습니다.

(i, j) ~ (x, y)의 합 구하기

위에서 구한 합을 활용해서 (i, j) ~ (x, y)의 합을 구할 수 있습니다. 우리가 구하고자 하는 부분은 빨간 빗금이 쳐진 (i, j) ~ (x, y)의 합입니다.

먼저 전체 사각형에 해당하는 dp(x, y)에서 각각 파란색과 초록색 빗금에 해당하는 dp(i - 1, y)와 dp(x, j - 1)을 빼줍니다. 그리고 두 번 빼진 dp(i - 1, j - 1)을 더하면 됩니다.

// 입력 받기
let input = readLine()!.split(separator: " ").map { Int(String($0))! }
let N = input[0], M = input[1]

// 이차원 배열 및 cache 이차원 배열 만들기
    //👉 0행과 0열에 0으로 한 줄씩 추가해서 예외처리를 편하게 할 수 있도록 (0은 합에 영향 없음)
var matrix = [[Int]]()
matrix.append(Array(repeating: 0, count: M + 1))
    // 0 * (M + 1)의 행을 추가
var cache: [[Int?]] = Array(repeating: Array(repeating: nil, count: M + 1), count: N + 1)
for _ in 0..<N {
    matrix.append([0] + readLine()!.split(separator: " ").map { Int(String($0))! })
    //각 행의 맨 앞에 0 추가 (0열을 모두 0으로)
}

// (0, 0) ~ (r, c)의 합을 구하는 함수
func dp(_ r: Int, _ c: Int) -> Int {
    // 0행 혹은 0열이면 0
    if r == 0 || c == 0 {
        cache[r][c] = 0
    }
    
    // 초기값 세팅 (1, 1) ~ (1, 1)의 합은 matrix[1][1]
    if r == 1 && c == 1 {
        cache[r][c] = matrix[r][c]
    }
    
    // 점화식
    if cache[r][c] == nil {
        cache[r][c] = dp(r - 1, c) + dp(r, c - 1) - dp(r - 1, c - 1) + matrix[r][c]
    }
    
    return cache[r][c]!
}

// 각 case 별로 합 구하기
let T = Int(readLine()!)!
for _ in 0..<T {
    let line = readLine()!.split(separator: " ").map { Int(String($0))! }
    let i = line[0], j = line[1], x = line[2], y = line[3]
    print(dp(x, y) - dp(x, j - 1) - dp(i - 1, y) + dp(i - 1, j - 1))
}