3294 - Hmmm

Cerința

Fie λ o permutare de grad N și K un număr natural nenul. Să se afișeze toate soluțiile ecuației x^K=λ în ordine lexicografică.

Date de intrare

Fișierul de intrare hmmm.in conține pe prima linie gradul permutării N și K, iar pe a doua linie se citește permutarea λ.

Date de ieșire

Fișierul de ieșire hmmm.out va conține toate soluțiile x ale ecuației în ordine lexicografică, câte una pe linie. Elementele permutărilor se separă printr-un spațiu.

Restricții și precizări

• N ≤ 9
• 2 ≤ K ≤ 1.000.000.000
• Întotdeauna există cel puțin o soluție.
• Pentru teste în valoare de 50 de puncte K ≤ 15

Exemplu 1

Intrare

4 2
1 2 3 4

Iesire

1 2 3 4
1 2 4 3
1 3 2 4
1 4 3 2
2 1 3 4
2 1 4 3
3 2 1 4
3 4 1 2
4 2 3 1
4 3 2 1

Rezolvare

import itertools


def read_input():
    with open("hmmm.in", "r") as file:
        N, K = map(int, file.readline().strip().split())
        lambda_permutation = list(map(int, file.readline().strip().split()))
    return N, K, lambda_permutation


def write_output(solutions):
    with open("hmmm.out", "w") as file:
        for solution in solutions:
            file.write(" ".join(map(str, solution)) + "\n")


def apply_permutation(permutation, k):
    n = len(permutation)
    result = list(range(1, n + 1))
    for _ in range(k):
        result = [result[permutation[i] - 1] for i in range(n)]
    return result


def generate_solutions(N, K, lambda_permutation):
    all_permutations = itertools.permutations(range(1, N + 1))
    solutions = []

    for permutation in all_permutations:
        if apply_permutation(list(permutation), K) == lambda_permutation:
            solutions.append(permutation)

    solutions.sort()
    return solutions


def main():
    N, K, lambda_permutation = read_input()

    # Validări
    if not (1 <= N <= 9):
        raise ValueError("N trebuie să fie între 1 și 9")
    if not (2 <= K <= 1000000000):
        raise ValueError("K trebuie să fie între 2 și 1.000.000.000")

    solutions = generate_solutions(N, K, lambda_permutation)
    write_output(solutions)


if __name__ == "__main__":
    main()