1950 - PXP

Cerinţa

Se dă un şir format din N numere naturale nenule. Spunem că un număr e fericit dacă se poate scrie ca suma pătratelor a două numere naturale. Notăm cu K numărul numerelor fericite din şir şi cu P produsul acestora. Aflaţi numărul K precum şi două numere naturale care au suma pătratelor egală cu PE, unde E este un număr natural dat.

Date de intrare

Fișierul de intrare pxp.in conține pe prima linie numerele N şi E , iar pe a doua linie N numere naturale separate prin spațiu.

Date de ieșire

Fișierul de ieșire pxp.out va conține pe prima linie numărul K, reprezentând numărul numerelor fericite din şir iar pe a doua linie cele două numere care au suma pătratelor egală cu PE.

Restricţii şi precizări

• 1 ≤ N ≤ 100.000
• 1 ≤ E ≤ 3
• numerele din şir sunt mai mici decât 1.000.001 şi cel puţin un număr din şir este fericit
• 1 ≤ P ≤ 1036
• pentru a doua cerinţă se va afişa orice soluţie, fiecare test având soluţie
• pentru prima cerinţă se acordă 40p, iar pentru a doua 60p
• 0% din teste au P ≤ 1000000, iar alte 40% au P ≤ 1018

Exemplul 1

pxp.in

5 1 11 2 13 7 5

pxp.out

3 3 11

Explicație

Avem N=5 şi E=1. În şir sunt 3 numere fericite, 2, 13 şi 5, deoarece 2=12+12 , 13=22+32 şi 5=12+22. Produsul lor este 130=32+112, deci două numere pentru care suma pătratelor este P1 sunt 3 şi 11.

Rezolvare

def is_happy_number(num):
    i = 1
    while i * i <= num:
        if num - i * i == int((num - i * i) ** 0.5) ** 2:
            return True
        i += 1
    return False

with open("pxp.in", "r") as fin:
    N, E = map(int, fin.readline().split())
    numbers = list(map(int, fin.readline().split()))

happy_numbers = []
product = 1
for num in numbers:
    if is_happy_number(num):
        happy_numbers.append(num)
        product *= num

K = len(happy_numbers)

target_sum = product * E
found = False
for i in range(len(happy_numbers)):
    for j in range(i + 1, len(happy_numbers)):
        if happy_numbers[i] ** 2 + happy_numbers[j] ** 2 == target_sum:
            result = (happy_numbers[i], happy_numbers[j])
            found = True
            break
    if found:
        break

with open("pxp.out", "w") as fout:
    fout.write(str(K) + "\n")
    fout.write(" ".join(map(str, result)))