Generalization of the 'Fitch' Cheney card trick

cheney.py

#! /usr/bin/env python3

''' Generalization of the 'Fitch' Cheney card trick

    Written by PM 2Ring 2003.02.13
    Converted to Python 2018.06.27
'''

import sys
import readline
from math import factorial
from itertools import combinations

def unrank(seq, perm_index):
    ''' Permute seq by lexicographic perm_index '''
    # Convert the permutation index to factorial base form.
    fcode = [0]
    for i in range(2, len(seq) + 1):
        perm_index, r = divmod(perm_index, i)
        fcode.append(r)

    # Build the permutation from the factorial base code
    seq = seq[:]
    return [seq.pop(u) for u in reversed(fcode)]

def rank(seq):
    ''' Determine the lexicographic permutation index of seq,
        which cannot contain repeated items.
    '''
    # Find the permutation index in factorial base form
    fcode = [sum(u < v for u in seq[i:])
        for i, v in enumerate(seq[:-1], 1)]

    #Convert the factorial base code to integer
    perm_index, base = 0, 1
    for j, v in enumerate(reversed(fcode), 2):
        perm_index += v * base
        base *= j
    return perm_index

def cheney_encode(cards, num):
    cards = sorted(cards)

    # Calculate the checksum to select the index of the hidden card
    idx = sum(cards) % num
    hidden = cards[idx]

    # Calculate the required permutation index
    perm_index = (hidden - idx) // num
    del cards[idx]

    # Permute the shown cards
    return unrank(cards, perm_index), hidden

def cheney_decode(cards, num):
    # Find the permutation index
    perm_index = rank(cards)

    # Calculate the negative checksum
    checksum = -sum(cards) % num

    # Calculate the value of the hidden card
    value = num * perm_index + checksum

    # Adjust the value to skip over shown cards
    for u in sorted(cards):
        if u <= value:
            value += 1
    return value

def cheney_test(num):
    decksize = factorial(num) + num - 1
    deck = range(decksize)
    print('Deck', deck)
    for i, cards in enumerate(combinations(deck, num), 1):
        if i % 10000 == 0: 
            print(f'  {i}', end='\r') 
        shown, hidden = cheney_encode(cards, num)
        found = cheney_decode(shown, num)
        if found != hidden:
            print(i, cards, shown, hidden, found)
    print()

def permute_test(num):
    seq = list(range(num))
    for i in range(factorial(num)):
        permuted = unrank(seq, i)
        j = rank(permuted)
        print(f'{i:3}: {permuted} {j} {i==j}')

def main(num):
    decksize = factorial(num) + num - 1
    deck = set(range(decksize))
    print('The "Fitch" Cheney card trick, integer version.\n'
        f'Deck size = {decksize}\n'
        f'Enter {num} numbers to encode or {num-1} numbers to decode.\n'
        'All numbers must be unique, and must be equal to or greater '
        f'than 0 and less than {decksize}.\n'
        'Separate numbers using spaces. Hit Ctrl-C to exit.')
    while True:
        try:
            src = input('> ')
            cards = [int(u) for u in src.split()]
            if not cards:
                continue
            scards = set(cards)
            len_cards = len(cards)
            if len_cards != len(scards):
                raise ValueError('No duplicates!')
            bad = scards - deck
            if bad:
                raise ValueError(f'Bad values: {bad}')
            if len_cards == num:
                shown, hidden = cheney_encode(cards, num)
                print(' '.join([str(u) for u in shown]), ':', hidden)
            elif len_cards == num - 1:
                print(cheney_decode(cards, num))
            else:
                raise ValueError('Bad quantity') 
        except ValueError as err:
            print(err)
        except KeyboardInterrupt:
            print('\nBye')
            break

if __name__ == '__main__':
    num = int(sys.argv[1]) if len(sys.argv) > 1 else 4
    #cheney_test(num)
    #permute_test(num)
    main(num)