Merge pull request #249 from skirpichev/perm/225

Add perm()

Author: skirpichev

gmp.c

@@ -1911,6 +1911,72 @@ gmp_comb(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
     return NULL;
 }
 
+static PyObject *
+gmp_perm(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
+{
+    if (nargs > 2 || nargs < 1) {
+        PyErr_SetString(PyExc_TypeError, "one or two arguments required");
+        return NULL;
+    }
+    if (nargs == 1) {
+        return gmp_fac(self, args[0]);
+    }
+
+    MPZ_Object *x, *y, *res = MPZ_new(0);
+
+    if (!res) {
+        return NULL; /* LCOV_EXCL_LINE */
+    }
+    CHECK_OP_INT(x, args[0]);
+    CHECK_OP_INT(y, args[1]);
+    if (zz_isneg(&x->z) || zz_isneg(&y->z)) {
+        PyErr_SetString(PyExc_ValueError,
+                        "perm() not defined for negative values");
+        goto err;
+    }
+
+    int64_t n, k;
+
+    if ((zz_to_i64(&x->z, &n) || n > ULONG_MAX)
+        || (zz_to_i64(&y->z, &k) || k > ULONG_MAX))
+    {
+        PyErr_Format(PyExc_OverflowError,
+                     "perm() arguments should not exceed %ld",
+                     ULONG_MAX);
+        goto err;
+    }
+    Py_XDECREF(x);
+    Py_XDECREF(y);
+    if (k > n) {
+        return (PyObject *)res;
+    }
+
+    MPZ_Object *den = MPZ_new(0);
+
+    if (!den) {
+        /* LCOV_EXCL_START */
+        PyErr_NoMemory();
+        goto err;
+        /* LCOV_EXCL_STOP */
+    }
+    if (zz_fac((uint64_t)n, &res->z)
+        || zz_fac((uint64_t)(n-k), &den->z)
+        || zz_div(&res->z, &den->z, ZZ_RNDD, &res->z, NULL))
+    {
+        /* LCOV_EXCL_START */
+        Py_DECREF(den);
+        PyErr_NoMemory();
+        goto err;
+        /* LCOV_EXCL_STOP */
+    }
+    Py_DECREF(den);
+    return (PyObject *)res;
+err:
+end:
+    Py_DECREF(res);
+    return NULL;
+}
+
 static zz_rnd
 get_round_mode(PyObject *rndstr)
 {
@@ -2124,6 +2190,9 @@ static PyMethodDef gmp_functions[] = {
      ("comb($module, n, k, /)\n--\n\nNumber of ways to choose k"
       " items from n items without repetition and order.\n\n"
       "Also called the binomial coefficient.")},
+    {"perm", (PyCFunction)gmp_perm, METH_FASTCALL,
+     ("perm($module, n, k=None, /)\n--\n\nNumber of ways to choose k"
+      " items from n items without repetition and withorder.")},
     {"_mpmath_normalize", (PyCFunction)gmp__mpmath_normalize, METH_FASTCALL,
      NULL},
     {"_mpmath_create", (PyCFunction)gmp__mpmath_create, METH_FASTCALL, NULL},
@@ -2216,7 +2285,7 @@ gmp_exec(PyObject *m)
                        "numbers.Integral.register(gmp.mpz)\n"
                        "gmp.fac = gmp.factorial\n"
                        "gmp.__all__ = ['comb', 'factorial', 'gcd', 'isqrt',\n"
-                       "               'lcm', 'mpz']\n"
+                       "               'lcm', 'mpz', 'perm']\n"
                        "gmp.__version__ = imp.version('python-gmp')\n");
 
     PyObject *res = PyRun_String(str, Py_file_input, ns, ns);

tests/test_functions.py

@@ -16,6 +16,7 @@
     isqrt_rem,
     lcm,
     mpz,
+    perm,
 )
 from hypothesis import example, given
 from hypothesis.strategies import booleans, integers, lists, sampled_from
@@ -59,6 +60,21 @@ def test_comb(x, y):
     assert comb(x, y) == r
 
 
+@given(integers(min_value=0, max_value=12345),
+       integers(min_value=0, max_value=12345))
+def test_perm(x, y):
+    mx = mpz(x)
+    my = mpz(y)
+    r = math.perm(x, y)
+    assert perm(mx, my) == r
+    assert perm(mx, y) == r
+    assert perm(x, my) == r
+    assert perm(x, y) == r
+    rx = math.factorial(x)
+    assert perm(mx) == rx
+    assert perm(x) == rx
+
+
 @given(bigints(), bigints(), bigints())
 @example(1<<(67*2), 1<<65, 1)
 @example(123, 1<<70, 1)
@@ -187,6 +203,16 @@ def test_interfaces():
         comb(2**1000, 1)
     with pytest.raises(OverflowError):
         comb(1, 2**1000)
+    with pytest.raises(TypeError):
+        perm(1, 2, 3)
+    with pytest.raises(ValueError, match="not defined for negative values"):
+        perm(-1, 2)
+    with pytest.raises(ValueError, match="not defined for negative values"):
+        perm(2, -1)
+    with pytest.raises(OverflowError):
+        perm(2**1000, 1)
+    with pytest.raises(OverflowError):
+        perm(1, 2**1000)
     with pytest.raises(TypeError):
         _mpmath_create(1j)
     with pytest.raises(TypeError):