311. Sparse Matrix Multiplication

import java.util.HashMap;

/**
 * First convert the input matrix to sparse format (into HashMap of HashMaps).
 * Then calculate the product and the convert back the result to 2D matrix.
 *
 * Time Complexity: Convert A[][] to sparseA -> O(rA*cA). Convert B[][] to
 * sparseB -> O(rB*cB). Product -> O(rA*cA*cB). Convert sparseC to C[][] ->
 * O(rA*cB). Thus total time Complexity -> O(rA*cA*cB).
 *
 * Space Complexity: If input/output is in 2D array -> O(rA*cA + rB*cB + rA*cB).
 * If input/output is in sparse format -> O(1).
 *
 * In Facebook Dot Product of Sparse Vector is asked.
 * https://www.cnblogs.com/EdwardLiu/p/6399867.html
 */
class Solution {
    public int[][] multiply(int[][] A, int[][] B) throws IllegalArgumentException {
        if (A == null || B == null) {
            return new int[0][0];
        }
        if (A[0].length != B.length) {
            throw new IllegalArgumentException("Number of columns in A should be equal to number of rows in B");
        }

        int rowA = A.length;
        HashMap<Integer, HashMap<Integer, Integer>> sparseA = getSparseMatrix(A);

        int colB = B[0].length;
        HashMap<Integer, HashMap<Integer, Integer>> sparseB = getSparseMatrix(B);

        HashMap<Integer, HashMap<Integer, Integer>> sparseC = new HashMap<>();

        for (int xA : sparseA.keySet()) {
            for (int yA : sparseA.get(xA).keySet()) {
                if (!sparseB.containsKey(yA)) {
                    continue;
                }
                for (int yB : sparseB.get(yA).keySet()) {
                    int val = sparseA.get(xA).get(yA) * sparseB.get(yA).get(yB);
                    if (!sparseC.containsKey(xA)) {
                        sparseC.put(xA, new HashMap<>());
                    }
                    int preVal = sparseC.get(xA).getOrDefault(yB, 0);
                    sparseC.get(xA).put(yB, preVal + val);
                }
            }
        }

        int[][] C = new int[rowA][colB];
        for (int xC : sparseC.keySet()) {
            for (int yC : sparseC.get(xC).keySet()) {
                C[xC][yC] = sparseC.get(xC).get(yC);
            }
        }

        return C;
    }

    private HashMap<Integer, HashMap<Integer, Integer>> getSparseMatrix(int[][] M) {
        HashMap<Integer, HashMap<Integer, Integer>> sparseM = new HashMap<>();
        if (M == null || M.length == 0 || M[0].length == 0) {
            return sparseM;
        }

        for (int i = 0; i < M.length; i++) {
            for (int j = 0; j < M[0].length; j++) {
                if (M[i][j] != 0) {
                    if (!sparseM.containsKey(i)) {
                        sparseM.put(i, new HashMap<>());
                    }
                    sparseM.get(i).put(j, M[i][j]);
                }
            }
        }

        return sparseM;
    }
}