코딩테스트 연습 - 행렬의 곱셈

문제 풀이 아이디어

행렬의 곱셈의 원리를 코드로 구현해야 합니다. a x b 행렬과 b x c 행렬을 곱하는 경우 a x c 행렬이 됩니다. 그리고 각 첫 번째 행렬의 n행과 두 번째 행렬의 m열의 모든 원소를 순서대로 곱해서 더한 값이 결과 행렬의 n행 m열의 값이 됩니다.

코드로는 3중 반복문으로 구현할 수 있습니다. 먼저 arr1의 각 행, arr2의 각 열을 반복문으로 순회하면서 곱한 값을 구하면 됩니다. (2중 반복문) 그리고 곱한 값을 구하는 과정에서 반복문을 하나 더 사용해서 3중 반복문입니다.

시간복잡도는 O(N^3)이지만 행과 열의 길이가 최대 100이므로 시간초과를 걱정할 필요는 없습니다.

코드

반복문의 복잡도를 낮추기 위해서 arr2에서 열을 구하는 함수와 두 배열을 곱하는 함수를 별도로 선언해서 사용했습니다.

func solution(_ arr1:[[Int]], _ arr2:[[Int]]) -> [[Int]] {
    // arr2에서 n번째 열을 구하는 함수
    func col(_ n: Int) -> [Int] {
        var result = [Int]()
        for i in 0..<arr2.count {
            result.append(arr2[i][n])
        }
        return result
    }
    
    // 두 배열을 곱하는 함수
    func multiplyArray(_ a: [Int], _ b: [Int]) -> Int {
        var result = 0
        for i in 0..<a.count {
            result += a[i] * b[i]
        }
        return result
    }
    
    var ans = [[Int]]()
    
    // arr1의 행과 arr2의 열을 순회하면서 서로 곱함
    for i in 0..<arr1.count {
        // arr1의 n행과 각 열을 곱한 값의 배열이 결과의 n행이 됨
        var arr = [Int]()
        for j in 0..<arr2[0].count {
            let result = multiplyArray(arr1[i], col(j))
            arr.append(result)
        }
        ans.append(arr)
    }
    
    return ans
}