switch eval_monad_expand_where to use np.repeat

briangu · Dec 8, 2024 · 0fad0ea · 0fad0ea
1 parent 6710bbf
commit 0fad0ea
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/klongpy/monads.py b/klongpy/monads.py
@@ -74,7 +74,8 @@ def eval_monad_expand_where(a):
                   &[0 1 0 1 0]  -->   [1 3]
 
     """
-    return np.concatenate([np.zeros(x, dtype=int) + i for i,x in enumerate(a if is_list(a) else [a])])
+    arr = a if is_list(a) else [a]
+    return np.repeat(np.arange(len(arr)), arr)
 
 
 def eval_monad_first(a):