faster tril_indices

nonzero is slow. Use np.indices and np.broadcast_to to speed it up.

1 files changed, 8 insertions, 2 deletions

diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py
index 2b4cbdfbb..8e008ba71 100644
--- a/numpy/lib/twodim_base.py
+++ b/numpy/lib/twodim_base.py
@@ -894,7 +894,10 @@ def tril_indices(n, k=0, m=None):
            [-10, -10, -10, -10]])
 
     """
-    return nonzero(tri(n, m, k=k, dtype=bool))
+    tri = np.tri(n, m=m, k=k, dtype=bool)
+
+    return tuple(np.broadcast_to(inds, tri.shape)[tri]
+                 for inds in np.indices(tri.shape, sparse=True))
 
 
 def _trilu_indices_form_dispatcher(arr, k=None):
@@ -1010,7 +1013,10 @@ def triu_indices(n, k=0, m=None):
            [ 12,  13,  14,  -1]])
 
     """
-    return nonzero(~tri(n, m, k=k-1, dtype=bool))
+    tri = ~np.tri(n, m, k=k - 1, dtype=bool)
+
+    return tuple(np.broadcast_to(inds, tri.shape)[tri]
+                 for inds in np.indices(tri.shape, sparse=True))
 
 
 @array_function_dispatch(_trilu_indices_form_dispatcher)