Find lexicographical predecessor or successor (permutation)

lex_perm.py

import sys

### Test

def test_predecessor():
    """Basic unit tests for find_predecessor."""
    assert main('1432', 0) == '1423'
    assert main('04123', 0) == '03421'
    assert main('315426', 0) == '315264'
    assert main('315246', 0) == '314652'
    assert main('ABDC', 0) == 'ABCD'
    assert main('ABEZKFC', 0) == 'ABEZKCF'
    assert main('CZEFGHL', 0) == 'CLZHGFE'
    assert main('12345', 0) == '12345'

def test_successor():
    """Basic unit tests for find_predecessor."""
    assert main('1423', 1) == '1432'
    assert main('03421', 1) == '04123'
    assert main('315264', 1) == '315426'
    assert main('314652', 1) == '315246'
    assert main('ABCD', 1) == 'ABDC'
    assert main('ABEZKCF', 1) == 'ABEZKFC'
    assert main('CLZHGFE', 1) == 'CZEFGHL'
    assert main('54321', 1) == '54321'

### Helpers

def update_seq(seq, left, right):
    """Return updated sequence based on left and right indices.
    First swap left and right indices; then reverse subsequence after left index.
    """
    seq[left], seq[right] = seq[right], seq[left]
    return seq[:left+1] + seq[left+1:][::-1]

def find_min(seq, ind):
    """In subsequence right of ind, find index of min item that is > than seq[ind]."""
    items = [(ind + i, item) for i, item in enumerate(seq[ind:]) if item > seq[ind]]
    return min(items, key=lambda x: x[1])[0]

def find_max(seq, ind):
    """In subsequence right of ind, find index of max item that is < than seq[ind]."""
    items = [(ind + i, item) for i, item in enumerate(seq[ind:]) if item < seq[ind]]
    return max(items, key=lambda x: x[1])[0]

def compare(fst, snd, successor):
    """Return correct comparison based on whether successor arg is True."""
    return snd < fst if successor else fst < snd

### Permutation

def find_perm(seq, successor):
    """Return lexicographical permutation that's predecessor or successor of sequence.
    Args
        seq: list, sequence to permute
        successor: bool, find successor if True, else find predecessor
    Returns
        Permutated sequence as a list.
    """

    find_item = find_min if successor else find_max 

    for ind in xrange(-1, -1 * len(seq), -1):
        next_ind = ind - 1

        if compare(seq[ind], seq[next_ind], successor):
            swap = find_item(seq, next_ind)
            seq = update_seq(seq, next_ind, swap)
            break
    
    return seq

### Main

def main(seq, successor=False):
    """Find perm that is lexicographical predecessor of sequence passed as string."""

    assert isinstance(seq, str), 'Expecting string as input sequence'
    ret = ''.join(find_perm(list(seq), successor))

    print seq, '\t->\t', ret
    return ret

if __name__ == '__main__':
    if len(sys.argv) == 2:
        main(sys.argv[1])
    elif len(sys.argv) == 3:
        main(sys.argv[1], int(sys.argv[2]))
    else:
        print 'Call with single arg for predecessor, eg $ python lex_perm.py ABDC'
        print 'A second arg of 0 is implied for predecessor; set to 1 for successor'