ENH: speed-up mlab.contiguous_regions using numpy

jaimefrio · jaimefrio · commit 6007164a23f7 · 2015-02-27T09:55:22.000-08:00
diff --git a/lib/matplotlib/mlab.py b/lib/matplotlib/mlab.py
@@ -3874,23 +3874,17 @@ def contiguous_regions(mask):
     """
     return a list of (ind0, ind1) such that mask[ind0:ind1].all() is
     True and we cover all such regions
-
-    TODO: this is a pure python implementation which probably has a much
-    faster numpy impl
     """
+    mask = np.asarray(mask, dtype=bool)
+
+    if not mask.size:
+        return []
 
-    in_region = None
-    boundaries = []
-    for i, val in enumerate(mask):
-        if in_region is None and val:
-            in_region = i
-        elif in_region is not None and not val:
-            boundaries.append((in_region, i))
-            in_region = None
+    idx, = np.nonzero(mask[:-1] != mask[1:])
+    idx = np.concatenate((([0],) if mask[0] else ()) + (idx + 1,) +
+                          (([len(mask)],) if mask[-1] else ()))
 
-    if in_region is not None:
-        boundaries.append((in_region, i+1))
-    return boundaries
+    return list(zip(idx[::2], idx[1::2]))
 
 
 def cross_from_below(x, threshold):
