ENH: use smallest integer for tri ranges

reduces buffer copy and comparison overhead for boolean outer product

1 files changed, 17 insertions, 2 deletions

diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py
index 5a0c0e7ee..20b5cdd67 100644
--- a/numpy/lib/twodim_base.py
+++ b/numpy/lib/twodim_base.py
@@ -11,10 +11,24 @@ __all__ = ['diag', 'diagflat', 'eye', 'fliplr', 'flipud', 'rot90', 'tri',
 
 from numpy.core.numeric import (
     asanyarray, subtract, arange, zeros, greater_equal, multiply, ones,
-    asarray, where, dtype as np_dtype, less
+    asarray, where, dtype as np_dtype, less, int8, int16, int32, int64
     )
+from numpy.core import iinfo
 
 
+i1 = iinfo(int8)
+i2 = iinfo(int16)
+i4 = iinfo(int32)
+def _min_int(low, high):
+    """ get small int that fits the range """
+    if high <= i1.max and low >= i1.min:
+        return int8
+    if high <= i2.max and low >= i2.min:
+        return int16
+    if high <= i4.max and low >= i4.min:
+        return int32
+    return int64
+
 
 def fliplr(m):
     """
@@ -396,7 +410,8 @@ def tri(N, M=None, k=0, dtype=float):
     if M is None:
         M = N
 
-    m = greater_equal.outer(arange(N), arange(-k, M-k))
+    m = greater_equal.outer(arange(N, dtype=_min_int(0, N)),
+                            arange(-k, M-k, dtype=_min_int(-k, M - k)))
 
     # Avoid making a copy if the requested type is already bool
     if np_dtype(dtype) != np_dtype(bool):