Question
Formatted question description: https://leetcode.ca/all/311.html
311 Sparse Matrix Multiplication
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:
A = [
[ 1, 0, 0],
[-1, 0, 3]
]
B = [
[ 7, 0, 0 ],
[ 0, 0, 0 ],
[ 0, 0, 1 ]
]
| 1 0 0 | | 7 0 0 | | 7 0 0 |
AB = | -1 0 3 | x | 0 0 0 | = | -7 0 3 |
| 0 0 1 |
Algorithm
Make sure that A[i][k]
is not 0 before continuing to calculate, and then traverse the kth row of matrix B. If B[K][J]
is not 0, accumulate the result matrix res[i][j] += A[i][k] * B[k][j]
, so that we can efficiently calculate the multiplication of the sparse matrix.
Code
Java
-
public class Sparse_Matrix_Multiplication { public int[][] multiply(int[][] A, int[][] B) { //validity check int[][] C = new int[A.length][B[0].length]; for (int i = 0; i < C.length; i++) { for (int k = 0; k < A[0].length; k++) { if (A[i][k] != 0) { // @note: non-zero check. if zero, then skip since result will be 0 for (int j = 0; j < C[0].length; j++) { C[i][j] += A[i][k] * B[k][j]; } } } } return C; } }
-
// OJ: https://leetcode.com/problems/sparse-matrix-multiplication/ // Time: O(MNK) // Space: O(1) class Solution { public: vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) { int M = A.size(), K = A[0].size(), N = B[0].size(); vector<vector<int>> ans(M, vector<int>(N)); for (int i = 0; i < M; ++i) { for (int j = 0; j < N; ++j) { for (int k = 0; k < K; ++k) { ans[i][j] += A[i][k] * B[k][j]; } } } return ans; } };
-
class Solution(object): def multiply(self, A, B): """ :type A: List[List[int]] :type B: List[List[int]] :rtype: List[List[int]] """ ret = [[0 for j in range(len(B[0]))] for i in range(len(A))] for i, row in enumerate(A): for k, a in enumerate(row): if a: for j, b in enumerate(B[k]): if b: ret[i][j] += a * b return ret