summaryrefslogtreecommitdiff
path: root/numpy/lib
diff options
context:
space:
mode:
authorDavid Warde-Farley <wardefar@iro.umontreal.ca>2011-01-07 13:36:18 -0500
committerCharles Harris <charlesr.harris@gmail.com>2011-01-10 20:30:44 -0700
commitf72c60510a225f242a36cdbbc3aacf1b36f05f22 (patch)
tree99e9aaaeefae2e854edb5bb0b9f92d0af969ae8c /numpy/lib
parentd43668c6284bc5cef4da155dcca6800b003294b7 (diff)
downloadnumpy-f72c60510a225f242a36cdbbc3aacf1b36f05f22.tar.gz
ENH: Add minlength keyword to bincount. Patch from ticket #1595.
Diffstat (limited to 'numpy/lib')
-rw-r--r--numpy/lib/src/_compiled_base.c32
-rw-r--r--numpy/lib/tests/test_function_base.py15
2 files changed, 39 insertions, 8 deletions
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):