diff options
Diffstat (limited to 'numpy/lib/twodim_base.py')
| -rw-r--r-- | numpy/lib/twodim_base.py | 10 |
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) |