diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/add_newdocs.py | 11 | ||||
| -rw-r--r-- | numpy/lib/src/_compiled_base.c | 32 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 15 |
3 files changed, 47 insertions, 11 deletions
diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py index baf5285f0..d7a8569fa 100644 --- a/numpy/add_newdocs.py +++ b/numpy/add_newdocs.py @@ -3796,12 +3796,15 @@ add_newdoc('numpy.lib._compiled_base', 'digitize', add_newdoc('numpy.lib._compiled_base', 'bincount', """ - bincount(x, weights=None) + bincount(x, weights=None, minlength=None) Count number of occurrences of each value in array of non-negative ints. The number of bins (of size 1) is one larger than the largest value in - `x`. Each bin gives the number of occurrences of its index value in `x`. + `x`. If `minlength` is specified, there will be at least this number + of bins in the output array (though it will be longer if necessary, + depending on the contents of `x`). + Each bin gives the number of occurrences of its index value in `x`. If `weights` is specified the input array is weighted by it, i.e. if a value ``n`` is found at position ``i``, ``out[n] += weight[i]`` instead of ``out[n] += 1``. @@ -3812,6 +3815,8 @@ add_newdoc('numpy.lib._compiled_base', 'bincount', Input array. weights : array_like, optional Weights, array of the same shape as `x`. + minlength : integer, optional + A minimum number of bins for the output array. Returns ------- @@ -3823,7 +3828,7 @@ add_newdoc('numpy.lib._compiled_base', 'bincount', ------ ValueError If the input is not 1-dimensional, or contains elements with negative - values. + values, or if `minlength` is non-positive. TypeError If the type of the input is float or complex. diff --git a/numpy/lib/src/_compiled_base.c b/numpy/lib/src/_compiled_base.c index 466fa0202..627fc3286 100644 --- a/numpy/lib/src/_compiled_base.c +++ b/numpy/lib/src/_compiled_base.c @@ -90,27 +90,29 @@ mnx (intp *i , intp len) /* * arr_bincount is registered as bincount. * - * bincount accepts one or two arguments. The first is an array of - * non-negative integers and the second, if present, is an array of weights, - * which must be promotable to double. Call these arguments list and + * bincount accepts one, two or three arguments. The first is an array of + * non-negative integers The second, if present, is an array of weights, + * which must be promotable to double. Call these arguments list and * weight. Both must be one-dimensional with len(weight) == len(list). If * weight is not present then bincount(list)[i] is the number of occurrences * of i in list. If weight is present then bincount(self,list, weight)[i] * is the sum of all weight[j] where list [j] == i. Self is not used. + * The third argument, if present, is a minimum length desired for the + * output array. */ static PyObject * arr_bincount(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds) { PyArray_Descr *type; - PyObject *list = NULL, *weight=Py_None; + PyObject *list = NULL, *weight=Py_None, *mlength=Py_None; PyObject *lst=NULL, *ans=NULL, *wts=NULL; - intp *numbers, *ians, len , mxi, mni, ans_size; + intp *numbers, *ians, len , mxi, mni, ans_size, minlength; int i; double *weights , *dans; - static char *kwlist[] = {"list", "weights", NULL}; + static char *kwlist[] = {"list", "weights", "minlength", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", - kwlist, &list, &weight)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO", + kwlist, &list, &weight, &mlength)) { goto fail; } if (!(lst = PyArray_ContiguousFromAny(list, PyArray_INTP, 1, 1))) { @@ -131,6 +133,20 @@ arr_bincount(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds) goto fail; } ans_size = numbers [mxi] + 1; + if (mlength != Py_None) { + if (!(minlength = PyArray_PyIntAsIntp(mlength))) { + goto fail; + } + if (minlength <= 0) { + /* superfluous, but may catch incorrect usage */ + PyErr_SetString(PyExc_ValueError, + "minlength must be positive"); + goto fail; + } + if (ans_size < minlength) { + ans_size = minlength; + } + } type = PyArray_DescrFromType(PyArray_INTP); if (weight == Py_None) { if (!(ans = PyArray_Zeros(1, &ans_size, type, 0))) { diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 4df13e5f9..700798e76 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -1026,6 +1026,21 @@ class TestBincount(TestCase): y = np.bincount(x, w) assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1])) + def test_with_minlength(self): + x = np.array([0, 1, 0, 1, 1]) + y = np.bincount(x, minlength=3) + assert_array_equal(y, np.array([2, 3, 0])) + + def test_with_minlength_smaller_than_maxvalue(self): + x = np.array([0, 1, 1, 2, 2, 3, 3]) + y = np.bincount(x, minlength=2) + assert_array_equal(y, np.array([1, 2, 2, 2])) + + def test_with_minlength_and_weights(self): + x = np.array([1, 2, 4, 5, 2]) + w = np.array([0.2, 0.3, 0.5, 0.1, 0.2]) + y = np.bincount(x, w, 8) + assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1, 0, 0])) class TestInterp(TestCase): def test_exceptions(self): |