fork(2) download 
  1. import itertools
  2.  
  3. _pascal_triangle = {(0, 0): 1}
  4.  
  5. def choose(n, r):
  6.   """Binomial coefficient nCr.
  7.  
  8.   Implemented using a cache. O(1) on cache hit, up to O(n*r) on miss.
  9.   """
  10.   if r < 0 or r > n:
  11.     return 0
  12.   key = (n, r)
  13.   res = _pascal_triangle.get(key)
  14.   if res is not None:
  15.     return res
  16.   res = choose(n - 1, r - 1) + choose(n - 1, r)
  17.   _pascal_triangle[key] = res
  18.   return res
  19.  
  20. def first_naive(i, n, r):
  21.   """Find first element and index of first combination with that first element.
  22.  
  23.   Returns a tuple of value and index.
  24.  
  25.   Naive implementation using an O(n) loop, with one choose invocation per loop.
  26.  
  27.   Example: first_naive(8, 5, 3) returns (1, 6) because the combination with
  28.   index 8 is [1, 3, 4] so it starts with 1, and because the first combination
  29.   that starts with 1 is [1, 2, 3] which has index 6.
  30.   """
  31.   s1 = 0
  32.   for k in range(n):
  33.     s2 = s1 + choose(n - k - 1, r - 1)
  34.     if i < s2:
  35.       return k, s1
  36.     s1 = s2
  37.  
  38. def first_bisect(i, n, r):
  39.   """Find first element and index of first combination with that first element.
  40.  
  41.   Returns a tuple of value and index.
  42.  
  43.   Efficient implementation using O(log n) bisection steps, with one choose
  44.   invocation in each step.
  45.  
  46.   Example: first_bisect(8, 5, 3) returns (1, 6) because the combination with
  47.   index 8 is [1, 3, 4] so it starts with 1, and because the first combination
  48.   that starts with 1 is [1, 2, 3] which has index 6.
  49.   """
  50.   nCr = choose(n, r)
  51.   k1 = r - 1
  52.   s1 = nCr
  53.   k2 = n
  54.   s2 = 0
  55.   while k2 - k1 > 1:
  56.     k3 = (k1 + k2) // 2
  57.     s3 = nCr - choose(k3, r)
  58.     if s3 <= i:
  59.       k2, s2 = k3, s3
  60.     else:
  61.       k1, s1 = k3, s3
  62.   return n - k2, s2
  63.  
  64. def combination(i, n, r):
  65.   """Compute combination with a given index.
  66.  
  67.   Equivalent to list(itertools.combinations(range(n), r))[i].
  68.  
  69.   Each combination is represented as a tuple of ascending elements, and
  70.   combinations are ordered lexicograplically.
  71.  
  72.   Args:
  73.     i: zero-based index of the combination
  74.     n: number of possible values, will be taken from range(n)
  75.     r: number of elements in result list
  76.   """
  77.   if r == 0:
  78.     return []
  79.   k, ik = first_bisect(i, n, r)
  80.   return tuple([k] + [j + k + 1 for j in combination(i - ik, n - k - 1, r - 1)])
  81.  
  82. def show_all(n, r):
  83.   """Print all combinations for a given n and r, in order."""
  84.   print("n={} r={}".format(n, r))
  85.   for i in range(choose(n, r)):
  86.     print("  {:8d}: {}".format(i, combination(i, n, r)))
  87.   print("")
  88.  
  89. def test_first():
  90.   """Check that first)naive and first_bisect are equivalent."""
  91.   for n in range(2, 10):
  92.     for r in range(1, n):
  93.       for i in range(choose(n, r)):
  94.         want = first_naive(i, n, r)
  95.         got = first_bisect(i, n, r)
  96.         assert want == got, "i={} n={} r={}: {} != {}".format(i, n, r, want, got)
  97.  
  98. def test_combination():
  99.   """Check that combination and itertools.combinations are equivalent."""
  100.   for n in range(2, 10):
  101.     for r in range(1, n):
  102.       got = [combination(i, n, r) for i in range(choose(n, r))]
  103.       want = list(itertools.combinations(range(n), r))
  104.       assert want == got, "n={} r={}: {} != {}".format(n, r, want, got)
  105.  
  106. show_all(5, 3)
  107. test_first()
  108. test_combination()
  109.  
Success #stdin #stdout 0.04s 27712KB
comments (?)
stdin 
Standard input is empty
stdout 
n=5 r=3
         0: (0, 1, 2)
         1: (0, 1, 3)
         2: (0, 1, 4)
         3: (0, 2, 3)
         4: (0, 2, 4)
         5: (0, 3, 4)
         6: (1, 2, 3)
         7: (1, 2, 4)
         8: (1, 3, 4)
         9: (2, 3, 4)
https://ideone.com/MTOEJQ
language:
Python 3 (python  3.12)
created:
6 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!