diff --git a/sw/decoders/proximal.py b/sw/decoders/proximal.py
index d1c72f7..953c0eb 100644
--- a/sw/decoders/proximal.py
+++ b/sw/decoders/proximal.py
@@ -27,6 +27,17 @@ class ProximalDecoder:
         self._gamma = gamma
         self._eta = eta
 
+        self._A = []
+        self._B = []
+
+        for row in self._H:
+            A_k = np.argwhere(row == 1)
+            self._A.append(A_k[:, 0])
+
+        for column in self._H.T:
+            B_k = np.argwhere(column == 1)
+            self._B.append(B_k[:, 0])
+
     @staticmethod
     def _L_awgn(s: np.array, y: np.array) -> np.array:
         """Variation of the negative log-likelihood for the special case of AWGN noise.
@@ -44,19 +55,11 @@ class ProximalDecoder:
         # Calculate second term
 
         for k, x_k in enumerate(x):
-            # TODO: Perform this operation for each row simultaneously
-            B_k = np.argwhere(self._H[:, k] == 1)
-            B_k = B_k[:, 0]  # Get rid of one layer of arrays
-
             # TODO: Perform the summation with np.sum()
             sum_result = 0
-            for i in B_k:
-                # TODO: Perform this operation for each column simultaneously
-                A_i = np.argwhere(self._H[i] == 1)
-                A_i = A_i[:, 0]  # Get rid of one layer of arrays
-
+            for i in self._B[k]:
                 prod = 1
-                for j in A_i:
+                for j in self._A[i]:
                     prod *= x[j]
 
                 sum_result += prod**2 - prod