diff --git a/optimask/_optimask.py b/optimask/_optimask.py
index a7d0d69..a602c5a 100644
--- a/optimask/_optimask.py
+++ b/optimask/_optimask.py
@@ -153,6 +153,7 @@ def _preprocess(x):
         iy, ix = np.empty(m * n, dtype=np.uint32), np.empty(m * n, dtype=np.uint32)
         cols_index_mapper = -np.ones(n, dtype=np.int32)
         rows_with_nan = np.empty(m, dtype=np.uint32)
+        cols_with_nan = np.empty(n, dtype=np.uint32)
         n_rows_with_nan = 0
         n_cols_with_nan = 0
         cnt = 0
@@ -168,6 +169,7 @@ def _preprocess(x):
                     else:
                         ix[cnt] = n_cols_with_nan
                         cols_index_mapper[j] = n_cols_with_nan
+                        cols_with_nan[n_cols_with_nan] = j
                         n_cols_with_nan += 1
                     cnt += 1
 
@@ -177,9 +179,7 @@ def _preprocess(x):
 
         iy, ix = iy[:cnt], ix[:cnt]
         rows_with_nan = rows_with_nan[:n_rows_with_nan]
-        cols_with_nan = np.flatnonzero(cols_index_mapper >= 0)[
-            cols_index_mapper[cols_index_mapper >= 0].argsort()
-        ].astype(np.uint32)
+        cols_with_nan = cols_with_nan[:n_cols_with_nan]
         return iy, ix, rows_with_nan, cols_with_nan
 
     def _trial(self, k, rng, m_nan, n_nan, iy, ix, m, n):
diff --git a/tests/test_optimask.py b/tests/test_optimask.py
index 0a3fa6b..cb66bde 100644
--- a/tests/test_optimask.py
+++ b/tests/test_optimask.py
@@ -126,14 +126,14 @@ def test_speed(opti_mask_instance):
     print("\nVertical arrays")
     for _ in range(5):
         start = perf_counter()
-        opti_mask_instance.solve(X=x)
-        print(f"{1e3 * (perf_counter() - start):.2f}ms")
+        rows, cols = opti_mask_instance.solve(X=x)
+        print(f"{1e3 * (perf_counter() - start):.2f}ms rows={rows.size} cols={cols.size}")
     x = generate_random(m=1_000, n=100_000, ratio=0.02)
     print("Horizontal arrays")
     for _ in range(5):
         start = perf_counter()
-        opti_mask_instance.solve(X=x)
-        print(f"{1e3 * (perf_counter() - start):.2f}ms")
+        rows, cols = opti_mask_instance.solve(X=x)
+        print(f"{1e3 * (perf_counter() - start):.2f}ms rows={rows.size} cols={cols.size}")
 
 
 def test_large_arrays(opti_mask_instance):