311. Sparse Matrix Multiplication

Given twosparse 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 |

Thoughts:

1. Idea from a CMU lecture.: A sparse matrix can be represented as a sequence of rows, each of which is a sequence of (column-number, value) pairs of the nonzero values in the row.
2. Time Complexity Proposal: O(m*n + k*nB). Here k: number of non-empty elements in A. So in the worst case (dense matrix), it's O(m*n*nB)(from here)

Code:

class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int m = A.length, n = A[0].length, nB = B[0].length;
        int [][] res = new int [m][nB];
        List[] rowA = new List[m];

        for(int i = 0; i < m; i++){
            List<Integer> colVal = new ArrayList<>();
            for (int j = 0 ; j < n; j ++){
                if(A[i][j]!= 0){
                  colVal.add(j);
                  colVal.add(A[i][j]);
                } 
            }
            rowA[i] = colVal;
        }

        for(int i = 0; i < m; i++){
            List<Integer> colVal = rowA[i];
            for(int p = 0; p < colVal.size(); p+=2){
                int colA = colVal.get(p);
                int valA = colVal.get(p + 1);
                for(int j = 0; j < nB; j++){
                    int valB = B[colA][j];
                    res[i][j] += valA * valB;
                }
            }
        }

        return res;
    }
}

Code: improvements: Definition of matrix multiplication: No extra space required

class Solution {
    public int[][] multiply(int[][] A, int[][] B) {
        int m = A.length, n = A[0].length, nB = B[0].length;

        int[][] res = new int[m][nB];

        for(int i = 0; i < m; i++){
            for(int k = 0; k < n; k++){
                if(A[i][k] != 0){
                    for(int j = 0; j < nB; j++){
                            if(B[k][j] != 0)
                            res[i][j] += A[i][k] * B[k][j];
                    }
                }
            }
        }

        return res;
    }
}