8000 ENH: Add minlength keyword to bincount. Patch from ticket #1595. · numpy/numpy@f72c605 · GitHub
[go: up one dir, main page]
Skip to content

Commit f72c605

Browse files
dwfcharris
dwf
authored and
charris
committed
ENH: Add minlength keyword to bincount. Patch from ticket #1595.
1 parent d43668c commit f72c605

File tree

3 files changed

+47
-11
lines changed

3 files changed

+47
-11
lines changed

‎numpy/add_newdocs.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3796,12 +3796,15 @@
37963796

37973797
add_newdoc('numpy.lib._compiled_base', 'bincount',
37983798
"""
3799-
bincount(x, weights=None)
3799+
bincount(x, weights=None, minlength=None)
38003800
38013801
Count number of occurrences of each value in array of non-negative ints.
38023802
38033803
The number of bins (of size 1) is one larger than the largest value in
3804-
`x`. Each bin gives the number of occurrences of its index value in `x`.
3804+
`x`. If `minlength` is specified, there will be at least this number
3805+
of bins in the output array (though it will be longer if necessary,
3806+
depending on the contents of `x`).
3807+
Each bin gives the number of occurrences of its index value in `x`.
38053808
If `weights` is specified the input array is weighted by it, i.e. if a
38063809
value ``n`` is found at position ``i``, ``out[n] += weight[i]`` instead
38073810
of ``out[n] += 1``.
@@ -3812,6 +3815,8 @@
38123815
Input array.
38133816
weights : array_like, optional
38143817
Weights, array of the same shape as `x`.
3818+
minlength : integer, optional
3819+
A minimum number of bins for the output array.
38153820
38163821
Returns
38173822
-------
@@ -3823,7 +3828,7 @@
38233828
------
38243829
ValueError
38253830
If the input is not 1-dimensional, or contains elements with negative
3826-
values.
3831+
values, or if `minlength` is non-positive.
38273832
TypeError
38283833
If the type of the input is float or complex.
38293834

‎numpy/lib/src/_compiled_base.c

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -90,27 +90,29 @@ mnx (intp *i , intp len)
9090
/*
9191
* arr_bincount is registered as bincount.
9292
*
93-
* bincount accepts one or two arguments. The first is an array of
94-
* non-negative integers and the second, if present, is an array of weights,
95-
* which must be promotable to double. Call these arguments list and
93+
* bincount accepts one, two or three arguments. The first is an array of
94+
* non-negative integers The second, if present, is an array of weights,
95+
* which must be promotable to double. Call these arguments list and
9696
* weight. Both must be one-dimensional with len(weight) == len(list). If
9797
* weight is not present then bincount(list)[i] is the number of occurrences
9898
* of i in list. If weight is present then bincount(self,list, weight)[i]
9999
* is the sum of all weight[j] where list [j] == i. Self is not used.
100+
* The third argument, if present, is a minimum length desired for the
101+
* output array.
100102
*/
101103
static PyObject *
102104
arr_bincount(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
103105
{
104106
PyArray_Descr *type;
105-
PyObject *list = NULL, *weight=Py_None;
107+
PyObject *list = NULL, *weight=Py_None, *mlength=Py_None;
106108
PyObject *lst=NULL, *ans=NULL, *wts=NULL;
107-
intp *numbers, *ians, len , mxi, mni, ans_size;
109+
intp *numbers, *ians, len , mxi, mni, ans_size, minlength;
108110
int i;
109111
double *weights , *dans;
110-
static char *kwlist[] = {"list", "weights", NULL};
112+
static char *kwlist[] = {"list", "weights", "minlength", NULL};
111113

112-
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O",
113-
kwlist, &list, &weight)) {
114+
if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|OO",
115+
kwlist, &list, &weight, &mlength)) {
114116
goto fail;
115117
}
116118
if (!(lst = PyArray_ContiguousFromAny(list, PyArray_INTP, 1, 1))) {
@@ -131,6 +133,20 @@ arr_bincount(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds)
131133
goto fail;
132134
}
133135
ans_size = numbers [mxi] + 1;
136+
if (mlength != Py_None) {
137+
if (!(minlength = PyArray_PyIntAsIntp(mlength))) {
138+
goto fail;
139+
}
140+
if (minlength <= 0) {
141+
/* superfluous, but may catch incorrect usage */
142+
PyErr_SetString(PyExc_ValueError,
143+
"minlength must be positive");
144+
goto fail;
145+
}
146+
if (ans_size < minlength) {
147+
ans_size = minlength;
148+
}
149+
}
134150
type = PyArray_DescrFromType(PyArray_INTP);
135151
if (weight == Py_None) {
136152
if (!(ans = PyArray_Zeros(1, &ans_size, type, 0))) {

‎numpy/lib/tests/test_function_base.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,6 +1026,21 @@ def test_simple_weight2(self):
10261026
y = np.bincount(x, w)
10271027
assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1]))
10281028

1029+
def test_with_minlength(self):
1030+
x = np.array([0, 1, 0, 1, 1])
1031+
y = np.bincount(x, minlength=3)
1032+
assert_array_equal(y, np.array([2, 3, 0]))
1033+
1034+
def test_with_minlength_smaller_than_maxvalue(self):
1035+
x = np.array([0, 1, 1, 2, 2, 3, 3])
1036+
y = np.bincount(x, minlength=2)
1037+
assert_array_equal(y, np.array([1, 2, 2, 2]))
1038+
1039+
def test_with_minlength_and_weights(self):
1040+
x = np.array([1, 2, 4, 5, 2])
1041+
w = np.array([0.2, 0.3, 0.5, 0.1, 0.2])
1042+
y = np.bincount(x, w, 8)
1043+
assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1, 0, 0]))
10291044

10301045
class TestInterp(TestCase):
10311046
def test_exceptions(self):

0 commit comments

Comments
 (0)
0