diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/lib/arraysetops.py | 35 | ||||
| -rw-r--r-- | numpy/lib/function_base.py | 2 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 12 |
3 files changed, 25 insertions, 24 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index 7bd666029..c7449eb12 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -34,28 +34,17 @@ import numpy ## # 03.11.2005, c -def ediff1d( ar1, to_end = None, to_begin = None ): +def ediff1d(ary, to_end = None, to_begin = None): """Array difference with prefixed and/or appended value.""" - dar1 = ar1[1:] - ar1[:-1] - if to_end and to_begin: - shape = (ar1.shape[0] + 1,) + ar1.shape[1:] - ed = numpy.empty( shape, dtype = ar1.dtype ) - ed[0], ed[-1] = to_begin, to_end - ed[1:-1] = dar1 - elif to_end: - ed = numpy.empty( ar1.shape, dtype = ar1.dtype ) - ed[-1] = to_end - ed[:-1] = dar1 - elif to_begin: - ed = numpy.empty( ar1.shape, dtype = ar1.dtype ) - ed[0] = to_begin - ed[1:] = dar1 - else: - ed = dar1 - + ary = numpy.asarray(ary) + ed = ary[1:] - ary[:-1] + if to_begin is not None: + ed = numpy.insert(ed, 0, to_begin) + if to_end is not None: + ed = numpy.append(ed, to_end) + return ed - ## # 01.11.2005, c # 02.11.2005 @@ -95,11 +84,11 @@ def intersect1d_nu( ar1, ar2 ): # 01.11.2005, c def setxor1d( ar1, ar2 ): """Set exclusive-or of 1D arrays with unique elements.""" - aux = numpy.concatenate( (ar1, ar2 ) ) + aux = numpy.concatenate((ar1, ar2)) aux.sort() - flag = ediff1d( aux, to_end = 1, to_begin = 1 ) == 0 - flag2 = ediff1d( flag, 0 ) == 0 - return aux.compress( flag2 ) + flag = ediff1d(aux, to_end = 1, to_begin = 1) == 0 + flag2 = ediff1d(flag) == 0 + return aux.compress(flag2) ## # 03.11.2005, c diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 0ede094f7..5bcd3a443 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1286,7 +1286,7 @@ def insert(arr, obj, values, axis=None): newshape = list(arr.shape) if isinstance(obj, (int, long, integer)): if (obj < 0): obj += N - if (obj < 0 or obj >=N): + if (obj < 0 or obj > N or (obj == N and N != 0)): raise ValueError, "index (%d) out of range (0<=index<=%d) "\ "in dimension %d" % (obj, N, axis) newshape[axis] += 1; diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index 706564410..ae98884d6 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -6,6 +6,7 @@ from numpy.testing import * set_package_path() import numpy from numpy.lib.arraysetops import * +from numpy.lib.arraysetops import ediff1d restore_path() ################################################## @@ -68,6 +69,17 @@ class test_aso(NumpyTestCase): c = setxor1d( a, b ) assert_array_equal( c, ec ) + def check_ediff1d(self): + zero_elem = numpy.array([]) + one_elem = numpy.array([1]) + two_elem = numpy.array([1,2]) + + assert_array_equal([],ediff1d(zero_elem)) + assert_array_equal([0],ediff1d(zero_elem,to_begin=0)) + assert_array_equal([0],ediff1d(zero_elem,to_end=0)) + assert_array_equal([-1,0],ediff1d(zero_elem,to_begin=-1,to_end=0)) + assert_array_equal([],ediff1d(one_elem)) + assert_array_equal([1],ediff1d(two_elem)) ## # 03.11.2005, c |