Leetcode #311：稀疏矩阵乘法

题目

题干

该问题稀疏矩阵乘法 题面：

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

给你两个 稀疏矩阵 A 和 B，请你返回 AB 的结果。
你可以默认 A 的列数等于 B 的行数。

示例

输入：
A = [[1,0,0],[-1,0,3]]
B = [[7,0,0],[0,0,0],[0,0,1]]
输出：
[[7,0,0], [-7,0,3]]

题解

由于是稀疏数组可以只计算非零数字。

将A矩阵转存为hashmap,然后对每一个非零数遍历B的列数来更新它在ans矩阵里可以被用到的位置。

C++

class Solution {
public:
    vector<vector<int>> multiply(vector<vector<int>>& A, vector<vector<int>>& B) {
        int n = A.size(), m = A[0].size(), k = B[0].size();
        vector<vector<int>> res(n, vector<int>(k, 0));
        unordered_map<int, int> am;
        //存储A的非零元素
        for (int i = 0; i < n; i ++)
            for (int j = 0; j < m; j ++)
            {
                if (!A[i][j]) continue;
                int idx = i * m + j;
                am[idx] = A[i][j];
            }
        for (auto aij : am)
        {
            int va = aij.second, aidx = aij.first, ai = aidx / m, aj = aidx % m;
            for (int j = 0; j < k; j ++)
            {
                res[ai][j] += va * B[aj][j];
            }
        }
        return res;
    }
};

Python

class Solution:
    def __init__(self,A,B) -> None:
        self.a = A
        self.b = B

    def multiply(self):
        am = {}
        n = len(self.a)
        m = len(self.a[0])
        k = len(self.b[0])
        res = []
        #存储A的非零元素
        for i in range(n):
            res = res+[[0]*k]
        for i in range(n):
            for j in range(m):
                if self.a[i][j] == 0:
                    continue
                idx = i * m + j
                am[idx] = self.a[i][j]
        for aidx, va in am.items():
            ai = aidx // m
            aj = aidx % m
            for j in range(k):
                res[ai][j] += va * self.b[aj][j]
        return res