311. Sparse Matrix Multiplication

https://leetcode.com/problems/sparse-matrix-multiplication/description/

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 |

Thoughts

实现矩阵乘法。res[i][j] = sigma_k(A[i][k] * B[k][j])，三重循环遍历i k j。

Code

class Solution {
public:
    vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
        vector<vector<int>> res(A[0].size(), vector<int>(B[0].size());
        for (int i = 0; i < A.size(); ++i) {
            for (int k = 0; k < A[0].size(); ++k) {
                if (A[i][k] != 0) {
                    for (int j = 0; j < B[0].size(); ++j) {
                        if (B[k][j] != 0) res[i][j] += A[i][k] * B[k][j];
                    }
                }
            }
        }
        return res;
    }
};

Analysis

当两个矩阵都是满的时候, 结果矩阵有M * Nb个元素, 计算每个元素需要花费O(N)时间, 因此时间复杂度O(MNbN), 空间O(MN + N * Nb).

Follow -up

[http://www.1point3acres.com/bbs/thread-201713-1-1.html] 如果给的是两个list分别代表两个vector, 每个list是根据index排好序的. 一种方法是依旧存到HashMap, 同样方法解 O(min(M, N)). 如果不让用hashmap, 就用两个指针从两个list头部开始往后遍历, 当list1的指向的index大就把list2的指针往前移动直到>=1的指针, 如果两个指针指向元素的index都相等, 就相乘加入结果. 时间复杂度O(M+N).