JIAKAOBO

LeetCode

venmo
wechat

感谢赞助!

  • ㊗️
  • 大家
  • offer
  • 多多！

Problem

Given two sparse matrices A and B, return the result of AB.

You may assume that A’s column number is equal to B’s row number.

Example:

Input:

A = [
  [ 1, 0, 0],
  [-1, 0, 3]
]

B = [
  [ 7, 0, 0 ],
  [ 0, 0, 0 ],
  [ 0, 0, 1 ]
]

Output:

     |  1 0 0 |   | 7 0 0 |   |  7 0 0 |
AB = | -1 0 3 | x | 0 0 0 | = | -7 0 3 |
                  | 0 0 1 |

Code

class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int rowA = A.length;
        int colA = A[0].length; // colA == rowB
        int colB = B[0].length;
        
        int[][] res = new int[rowA][colB];

        for(int i = 0; i < rowA; i++){
            for(int j = 0; j < colA; j++){
                if(A[i][j] != 0){
                    for(int k = 0; k < colB; k++){
                        if(B[j][k] != 0){
                            // res[i][k] = Aij * Bjk + Aij * Bjk + Aij * Bjk
                            // res[0][0] = A00 * B00 + A01 * B10 + A02 * B20
                            // i和k始终是0， j不停的变换0，1，2
                            res[i][k] += A[i][j] * B[j][k];
                        }
                    }
                }
            }
        }
        return res;
    }
}

class Solution:
    def multiply(self, A: List[List[int]], B: List[List[int]]) -> List[List[int]]:
        if len(A) == 0 or len(B) == 0:
            return [[]]

        a, b, c = len(A), len(B), len(B[0])
        res = [[0 for _ in range(c)] for _ in range(a)]

        for i in range(a):
            for j in range(b):
                if A[i][j] != 0:
                    for k in range(c):
                        if B[j][k] != 0:
                            res[i][k] += A[i][j] * B[j][k]

        return res