LeetCode 311. Sparse Matrix Multiplication

taoqick

于 2020-06-16 14:59:25 发布

阅读量219

点赞数

分类专栏：算法 leetcode

本文链接：https://blog.csdn.net/taoqick/article/details/106787392

版权

算法同时被 2 个专栏收录

474 篇文章 6 订阅

订阅专栏

leetcode

231 篇文章 0 订阅

订阅专栏

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 |

--------------------------------------------------------------------------

直接搞大家都会，比较好的方式是把B转换成列存储，列中非零的位置再去乘A。

class Solution:
    def multiply(self, A, B):
        cols = [[(j, b) for j, b in enumerate(col) if b!=0]
                for col in zip(*B)]
        return [[sum(row[j]*b for j, b in col)
                 for col in cols]
                 for row in A]