Combinatorics - permutations, variations, combinations

combinatorics.py

from enum import Enum


class OpType(Enum):
    permutation = 1
    variation = 4
    variation_repetition = 2
    combination = 5
    combination_repetition = 3


def permutations(input_set):
    return main_function(input_set, len(input_set), OpType.permutation)


def vartiations(input_set, k, repeat=False):
    ot = OpType.variation_repetition if repeat else OpType.variation
    return main_function(input_set, k, ot)


def combinations(input_set, k, repeat=False):
    ot = OpType.combination_repetition if repeat else OpType.combination
    return main_function(input_set, k, ot)


def main_function(input_set, k, operation_type):
    if k == 1:
        return [[a] for a in input_set]

    output = []
    temp_input = input_set[:]
    for x in input_set:
        temp = temp_input[:]
        if not (operation_type == OpType.combination_repetition
                or operation_type == OpType.variation_repetition):
            temp.remove(x)
        for var in main_function(temp, k - 1, operation_type):
            var.insert(0, x)
            output.insert(0, var)
        if operation_type == OpType.combination or operation_type == OpType.combination_repetition:
            temp_input.remove(x)
    return output

# print(permutations(['A', 'B', 'C']))
# print(vartiations(['A', 'B', 'C', 'D'], 2, False))
# print(combinations(['A', 'B', 'C', 'D'], 2, False))
# print(vartiations(['A', 'B', 'C', 'D'], 2, True))
# print(combinations(['A', 'B', 'C', 'D'], 2, True))