support_enumeration: Use util.combinatorics.next_k_array

instead of _next_k_combination and _next_k_array

quantecon/game_theory/support_enumeration.py

@@ -14,6 +14,7 @@
 import numpy as np
 from numba import jit
 from ..util.numba import _numba_linalg_solve
+from ..util.combinatorics import next_k_array
 
 
 def support_enumeration(g):
@@ -106,8 +107,8 @@ def _support_enumeration_gen(payoff_matrix0, payoff_matrix1):
                                                                actions)):
                             out[p][supp] = action[:-1]
                         yield out
-                _next_k_array(supps[1])
-            _next_k_array(supps[0])
+                next_k_array(supps[1])
+            next_k_array(supps[0])
 
 
 @jit(nopython=True, cache=True)
@@ -180,84 +181,3 @@ def _indiff_mixed_action(payoff_matrix, own_supp, opp_supp, A, out):
             if payoff > val:
                 return False
     return True
-
-
-@jit(nopython=True, cache=True)
-def _next_k_combination(x):
-    """
-    Find the next k-combination, as described by an integer in binary
-    representation with the k set bits, by "Gosper's hack".
-
-    Copy-paste from en.wikipedia.org/wiki/Combinatorial_number_system
-
-    Parameters
-    ----------
-    x : int
-        Integer with k set bits.
-
-    Returns
-    -------
-    int
-        Smallest integer > x with k set bits.
-
-    """
-    u = x & -x
-    v = u + x
-    return v + (((v ^ x) // u) >> 2)
-
-
-@jit(nopython=True, cache=True)
-def _next_k_array(a):
-    """
-    Given an array `a` of k distinct nonnegative integers, return the
-    next k-array in lexicographic ordering of the descending sequences
-    of the elements. `a` is modified in place.
-
-    Parameters
-    ----------
-    a : ndarray(int, ndim=1)
-        Array of length k.
-
-    Returns
-    -------
-    a : ndarray(int, ndim=1)
-        View of `a`.
-
-    Examples
-    --------
-    Enumerate all the subsets with k elements of the set {0, ..., n-1}.
-
-    >>> n, k = 4, 2
-    >>> a = np.arange(k)
-    >>> while a[-1] < n:
-    ...     print(a)
-    ...     a = _next_k_array(a)
-    ...
-    [0 1]
-    [0 2]
-    [1 2]
-    [0 3]
-    [1 3]
-    [2 3]
-
-    """
-    k = len(a)
-    if k == 0:
-        return a
-
-    x = 0
-    for i in range(k):
-        x += (1 << a[i])
-
-    x = _next_k_combination(x)
-
-    pos = 0
-    for i in range(k):
-        while x & 1 == 0:
-            x = x >> 1
-            pos += 1
-        a[i] = pos
-        x = x >> 1
-        pos += 1
-
-    return a