Return the rank of a partition with k parts which sums up to n, and also convert back

partition_rank_unrank.py

from functools import lru_cache


def p_single(n, k=0):
  k = k or n
  '''Return total number of partition for integer N of length k'''
  dp = [[0 for j in range(k+1)] for i in range(n+1)]
  for i in range(n+1):
    for j in range(k+1):
      if i == 0 and j == i:
        dp[i][j] = 1
      elif j == 1 or i == j:
        dp[i][j] = 1
      elif j > i:
        dp[i][j] = 0
      else:
        dp[i][j] = dp[i-1][j-1] + dp[i-j][j]
  return dp[n][k]


def p(n, k):
  '''Return total number of partition for integer N with length 1 up to k'''
  table = [[0]*(k+1) for i in range(n+1)]
  for i in range(k+1):
    table[0][i] = 1
  for i in range(1, n+1):
    table[i][1] = 1
  for i in range(1, n+1):
    for j in range(2, k+1):
      if i < j:
        table[i][j] = table[i][j-1]
      else:
        table[i][j] = table[i][j-1] + table[i-j][j]
  return table[n][k]


@lru_cache(maxsize=None)
def num_partitions(n, k, p):
  '''Returns the number of partitions of n with k parts and the largest part not exceeding p'''
  if n < 0 or k < 0 or p <= 0:
    return 0
  if n == 0 and k == 0:
    return 1
  return (num_partitions(n - p, k - 1, p)
          + num_partitions(n, k, p - 1))


def partition_rank(arr):
  '''Return the rank/index of a partition that has k parts and sums up to n'''
  n = _n = sum(arr)
  k = _k = len(arr)
  r = 0
  for v in arr:
    x = num_partitions(n, k, v-1)
    r += x
    n -= v
    k -= 1
  return p_single(_n, _k) - r - 1


def partition_unrank(n, k, index):
  '''Given the rank/index of a partition that has k parts and sums up to n
   return the original partition
  '''
  arr = [1] * k
  r = p_single(n, k) - index

  for i in range(k):
    for j in range(n+1):
      if num_partitions(n, k, j) >= r:
        arr[i] = j
        r -= num_partitions(n, k, j-1)
        n -= j
        k -= 1
        break
  return arr


# all partitions of 10 that have 3 parts
partitions = [(1, 1, 8), (1, 2, 7), (1, 3, 6), (2, 2, 6), (1, 4, 5), (2, 3, 5), (2, 4, 4), (3, 3, 4)]

for p in partitions:
  index = partition_rank(p)
  print(p, '->', index, '->', partition_unrank(10, 3, index))

# Output:
# (8, 1, 1) -> 0 -> [8, 1, 1]
# (7, 2, 1) -> 1 -> [7, 2, 1]
# (6, 3, 1) -> 2 -> [6, 3, 1]
# (6, 2, 2) -> 3 -> [6, 2, 2]
# (5, 4, 1) -> 4 -> [5, 4, 1]
# (5, 3, 2) -> 5 -> [5, 3, 2]
# (4, 4, 2) -> 6 -> [4, 4, 2]
# (4, 3, 3) -> 7 -> [4, 3, 3]