mirror of https://github.com/python/cpython
bpo-35431: Implemented math.comb (GH-11414)
This commit is contained in:
parent
5ac0b988fd
commit
4a686504eb
|
@ -232,6 +232,21 @@ Number-theoretic and representation functions
|
|||
:meth:`x.__trunc__() <object.__trunc__>`.
|
||||
|
||||
|
||||
.. function:: comb(n, k)
|
||||
|
||||
Return the number of ways to choose *k* items from *n* items without repetition
|
||||
and without order.
|
||||
|
||||
Also called the binomial coefficient. It is mathematically equal to the expression
|
||||
``n! / (k! (n - k)!)``. It is equivalent to the coefficient of k-th term in
|
||||
polynomial expansion of the expression ``(1 + x) ** n``.
|
||||
|
||||
Raises :exc:`TypeError` if the arguments not integers.
|
||||
Raises :exc:`ValueError` if the arguments are negative or if k > n.
|
||||
|
||||
.. versionadded:: 3.8
|
||||
|
||||
|
||||
Note that :func:`frexp` and :func:`modf` have a different call/return pattern
|
||||
than their C equivalents: they take a single argument and return a pair of
|
||||
values, rather than returning their second return value through an 'output
|
||||
|
|
|
@ -1862,6 +1862,57 @@ class IsCloseTests(unittest.TestCase):
|
|||
self.assertAllClose(fraction_examples, rel_tol=1e-8)
|
||||
self.assertAllNotClose(fraction_examples, rel_tol=1e-9)
|
||||
|
||||
def testComb(self):
|
||||
comb = math.comb
|
||||
factorial = math.factorial
|
||||
# Test if factorial defintion is satisfied
|
||||
for n in range(100):
|
||||
for k in range(n + 1):
|
||||
self.assertEqual(comb(n, k), factorial(n)
|
||||
// (factorial(k) * factorial(n - k)))
|
||||
|
||||
# Test for Pascal's identity
|
||||
for n in range(1, 100):
|
||||
for k in range(1, n):
|
||||
self.assertEqual(comb(n, k), comb(n - 1, k - 1) + comb(n - 1, k))
|
||||
|
||||
# Test corner cases
|
||||
for n in range(100):
|
||||
self.assertEqual(comb(n, 0), 1)
|
||||
self.assertEqual(comb(n, n), 1)
|
||||
|
||||
for n in range(1, 100):
|
||||
self.assertEqual(comb(n, 1), n)
|
||||
self.assertEqual(comb(n, n - 1), n)
|
||||
|
||||
# Test Symmetry
|
||||
for n in range(100):
|
||||
for k in range(n // 2):
|
||||
self.assertEqual(comb(n, k), comb(n, n - k))
|
||||
|
||||
# Raises TypeError if any argument is non-integer or argument count is
|
||||
# not 2
|
||||
self.assertRaises(TypeError, comb, 10, 1.0)
|
||||
self.assertRaises(TypeError, comb, 10, "1")
|
||||
self.assertRaises(TypeError, comb, "10", 1)
|
||||
self.assertRaises(TypeError, comb, 10.0, 1)
|
||||
|
||||
self.assertRaises(TypeError, comb, 10)
|
||||
self.assertRaises(TypeError, comb, 10, 1, 3)
|
||||
self.assertRaises(TypeError, comb)
|
||||
|
||||
# Raises Value error if not k or n are negative numbers
|
||||
self.assertRaises(ValueError, comb, -1, 1)
|
||||
self.assertRaises(ValueError, comb, -10*10, 1)
|
||||
self.assertRaises(ValueError, comb, 1, -1)
|
||||
self.assertRaises(ValueError, comb, 1, -10*10)
|
||||
|
||||
# Raises value error if k is greater than n
|
||||
self.assertRaises(ValueError, comb, 1, 10**10)
|
||||
self.assertRaises(ValueError, comb, 0, 1)
|
||||
|
||||
|
||||
|
||||
|
||||
def test_main():
|
||||
from doctest import DocFileSuite
|
||||
|
|
|
@ -0,0 +1,4 @@
|
|||
Implement :func:`math.comb` that returns binomial coefficient, that computes
|
||||
the number of ways to choose k items from n items without repetition and
|
||||
without order.
|
||||
Patch by Yash Aggarwal and Keller Fuchs.
|
|
@ -637,4 +637,53 @@ skip_optional_kwonly:
|
|||
exit:
|
||||
return return_value;
|
||||
}
|
||||
/*[clinic end generated code: output=aeed62f403b90199 input=a9049054013a1b77]*/
|
||||
|
||||
PyDoc_STRVAR(math_comb__doc__,
|
||||
"comb($module, /, n, k)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Number of ways to choose *k* items from *n* items without repetition and without order.\n"
|
||||
"\n"
|
||||
"Also called the binomial coefficient. It is mathematically equal to the expression\n"
|
||||
"n! / (k! * (n - k)!). It is equivalent to the coefficient of k-th term in\n"
|
||||
"polynomial expansion of the expression (1 + x)**n.\n"
|
||||
"\n"
|
||||
"Raises TypeError if the arguments are not integers.\n"
|
||||
"Raises ValueError if the arguments are negative or if k > n.");
|
||||
|
||||
#define MATH_COMB_METHODDEF \
|
||||
{"comb", (PyCFunction)(void(*)(void))math_comb, METH_FASTCALL|METH_KEYWORDS, math_comb__doc__},
|
||||
|
||||
static PyObject *
|
||||
math_comb_impl(PyObject *module, PyObject *n, PyObject *k);
|
||||
|
||||
static PyObject *
|
||||
math_comb(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
static const char * const _keywords[] = {"n", "k", NULL};
|
||||
static _PyArg_Parser _parser = {NULL, _keywords, "comb", 0};
|
||||
PyObject *argsbuf[2];
|
||||
PyObject *n;
|
||||
PyObject *k;
|
||||
|
||||
args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 2, 2, 0, argsbuf);
|
||||
if (!args) {
|
||||
goto exit;
|
||||
}
|
||||
if (!PyLong_Check(args[0])) {
|
||||
_PyArg_BadArgument("comb", 1, "int", args[0]);
|
||||
goto exit;
|
||||
}
|
||||
n = args[0];
|
||||
if (!PyLong_Check(args[1])) {
|
||||
_PyArg_BadArgument("comb", 2, "int", args[1]);
|
||||
goto exit;
|
||||
}
|
||||
k = args[1];
|
||||
return_value = math_comb_impl(module, n, k);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
/*[clinic end generated code: output=00aa76356759617a input=a9049054013a1b77]*/
|
||||
|
|
|
@ -2998,6 +2998,126 @@ math_prod_impl(PyObject *module, PyObject *iterable, PyObject *start)
|
|||
}
|
||||
|
||||
|
||||
/*[clinic input]
|
||||
math.comb
|
||||
|
||||
n: object(subclass_of='&PyLong_Type')
|
||||
k: object(subclass_of='&PyLong_Type')
|
||||
|
||||
Number of ways to choose *k* items from *n* items without repetition and without order.
|
||||
|
||||
Also called the binomial coefficient. It is mathematically equal to the expression
|
||||
n! / (k! * (n - k)!). It is equivalent to the coefficient of k-th term in
|
||||
polynomial expansion of the expression (1 + x)**n.
|
||||
|
||||
Raises TypeError if the arguments are not integers.
|
||||
Raises ValueError if the arguments are negative or if k > n.
|
||||
|
||||
[clinic start generated code]*/
|
||||
|
||||
static PyObject *
|
||||
math_comb_impl(PyObject *module, PyObject *n, PyObject *k)
|
||||
/*[clinic end generated code: output=bd2cec8d854f3493 input=565f340f98efb5b5]*/
|
||||
{
|
||||
PyObject *val = NULL,
|
||||
*temp_obj1 = NULL,
|
||||
*temp_obj2 = NULL,
|
||||
*dump_var = NULL;
|
||||
int overflow, cmp;
|
||||
long long i, terms;
|
||||
|
||||
cmp = PyObject_RichCompareBool(n, k, Py_LT);
|
||||
if (cmp < 0) {
|
||||
goto fail_comb;
|
||||
}
|
||||
else if (cmp > 0) {
|
||||
PyErr_Format(PyExc_ValueError,
|
||||
"n must be an integer greater than or equal to k");
|
||||
goto fail_comb;
|
||||
}
|
||||
|
||||
/* b = min(b, a - b) */
|
||||
dump_var = PyNumber_Subtract(n, k);
|
||||
if (dump_var == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
cmp = PyObject_RichCompareBool(k, dump_var, Py_GT);
|
||||
if (cmp < 0) {
|
||||
goto fail_comb;
|
||||
}
|
||||
else if (cmp > 0) {
|
||||
k = dump_var;
|
||||
dump_var = NULL;
|
||||
}
|
||||
else {
|
||||
Py_DECREF(dump_var);
|
||||
dump_var = NULL;
|
||||
}
|
||||
|
||||
terms = PyLong_AsLongLongAndOverflow(k, &overflow);
|
||||
if (terms < 0 && PyErr_Occurred()) {
|
||||
goto fail_comb;
|
||||
}
|
||||
else if (overflow > 0) {
|
||||
PyErr_Format(PyExc_OverflowError,
|
||||
"minimum(n - k, k) must not exceed %lld",
|
||||
LLONG_MAX);
|
||||
goto fail_comb;
|
||||
}
|
||||
else if (overflow < 0 || terms < 0) {
|
||||
PyErr_Format(PyExc_ValueError,
|
||||
"k must be a positive integer");
|
||||
goto fail_comb;
|
||||
}
|
||||
|
||||
if (terms == 0) {
|
||||
return PyNumber_Long(_PyLong_One);
|
||||
}
|
||||
|
||||
val = PyNumber_Long(n);
|
||||
for (i = 1; i < terms; ++i) {
|
||||
temp_obj1 = PyLong_FromSsize_t(i);
|
||||
if (temp_obj1 == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
temp_obj2 = PyNumber_Subtract(n, temp_obj1);
|
||||
if (temp_obj2 == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
dump_var = val;
|
||||
val = PyNumber_Multiply(val, temp_obj2);
|
||||
if (val == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
Py_DECREF(dump_var);
|
||||
dump_var = NULL;
|
||||
Py_DECREF(temp_obj2);
|
||||
temp_obj2 = PyLong_FromUnsignedLongLong((unsigned long long)(i + 1));
|
||||
if (temp_obj2 == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
dump_var = val;
|
||||
val = PyNumber_FloorDivide(val, temp_obj2);
|
||||
if (val == NULL) {
|
||||
goto fail_comb;
|
||||
}
|
||||
Py_DECREF(dump_var);
|
||||
Py_DECREF(temp_obj1);
|
||||
Py_DECREF(temp_obj2);
|
||||
}
|
||||
|
||||
return val;
|
||||
|
||||
fail_comb:
|
||||
Py_XDECREF(val);
|
||||
Py_XDECREF(dump_var);
|
||||
Py_XDECREF(temp_obj1);
|
||||
Py_XDECREF(temp_obj2);
|
||||
|
||||
return NULL;
|
||||
}
|
||||
|
||||
|
||||
static PyMethodDef math_methods[] = {
|
||||
{"acos", math_acos, METH_O, math_acos_doc},
|
||||
{"acosh", math_acosh, METH_O, math_acosh_doc},
|
||||
|
@ -3047,6 +3167,7 @@ static PyMethodDef math_methods[] = {
|
|||
{"tanh", math_tanh, METH_O, math_tanh_doc},
|
||||
MATH_TRUNC_METHODDEF
|
||||
MATH_PROD_METHODDEF
|
||||
MATH_COMB_METHODDEF
|
||||
{NULL, NULL} /* sentinel */
|
||||
};
|
||||
|
||||
|
|
Loading…
Reference in New Issue