bpo-35431: Implemented math.comb (GH-11414)

FR4NKESTI3N · rhettinger · commit 4a686504eb2b · 2019-06-01T00:21:27.000-07:00
diff --git a/Doc/library/math.rst b/Doc/library/math.rst
@@ -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
diff --git a/Lib/test/test_math.py b/Lib/test/test_math.py
@@ -1862,6 +1862,57 @@ def test_fractions(self):
         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
diff --git a/Misc/NEWS.d/next/Library/2019-01-02-19-48-23.bpo-35431.FhG6QA.rst b/Misc/NEWS.d/next/Library/2019-01-02-19-48-23.bpo-35431.FhG6QA.rst
@@ -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.
diff --git a/Modules/clinic/mathmodule.c.h b/Modules/clinic/mathmodule.c.h
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
@@ -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 */
 };
 
