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:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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 |

解法1：

只要存下非0的数即可，然后遍历每一对非0的pair，看是否col=row，然后加和到对应的单元格里。
C++

1

Java

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
public class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int row_a = A.length;
        int col_a = A[0].length;
        int row_b = B.length;
        int col_b = B[0].length;
        int[][] res = new int[row_a][col_b];    // store the result
        // res[i][j] = sum (A[i][k] .* B[k][j]) ...
        HashMap<Integer, Integer> a = new HashMap<Integer, Integer>();
        HashMap<Integer, Integer> b = new HashMap<Integer, Integer>();
        for (int i = 0; i < row_a; i++) {
            for (int j = 0; j < col_a; j++) {
                if (A[i][j] != 0) {
                    int index = i*col_a + j;
                    a.put(index, A[i][j]);
                }
            }
        }
        for (int i = 0; i < row_b; i++) {
            for (int j = 0; j < col_b; j++) {
                if (B[i][j] != 0) {
                    int index = i*col_b + j;
                    b.put(index, B[i][j]);
                }
            }
        }
        for (int index_a : a.keySet()) {
            for (int index_b : b.keySet()) {
                if (index_a % col_a == index_b / col_b) {
                    res[index_a / col_a][index_b % col_b] += a.get(index_a) * b.get(index_b);
                }
            }
        }
        return res;
    }
}