Merge pull request #7631 from perimosocordiae/patch-2

MAINT: linalg: fix comment, simplify math

1 files changed, 6 insertions, 6 deletions

diff --git a/numpy/linalg/linalg.py b/numpy/linalg/linalg.py
index f4acd0ef3..1f3ba7c54 100644
--- a/numpy/linalg/linalg.py
+++ b/numpy/linalg/linalg.py
@@ -2346,12 +2346,12 @@ def _multi_dot_three(A, B, C):
     than `_multi_dot_matrix_chain_order`
 
     """
-    # cost1 = cost((AB)C)
-    cost1 = (A.shape[0] * A.shape[1] * B.shape[1] +  # (AB)
-             A.shape[0] * B.shape[1] * C.shape[1])   # (--)C
-    # cost2 = cost((AB)C)
-    cost2 = (B.shape[0] * B.shape[1] * C.shape[1] +  #  (BC)
-             A.shape[0] * A.shape[1] * C.shape[1])   # A(--)
+    a0, a1b0 = A.shape
+    b1c0, c1 = C.shape
+    # cost1 = cost((AB)C) = a0*a1b0*b1c0 + a0*b1c0*c1
+    cost1 = a0 * b1c0 * (a1b0 + c1)
+    # cost2 = cost(A(BC)) = a1b0*b1c0*c1 + a0*a1b0*c1
+    cost2 = a1b0 * c1 * (a0 + b1c0)
 
     if cost1 < cost2:
         return dot(dot(A, B), C)