RosettaCodeData/Task/Ruth-Aaron-numbers/Python/ruth-aaron-numbers.py

105 lines
2.4 KiB
Python

from collections import deque
from itertools import count
from itertools import islice
from pprint import pprint
from typing import Iterable
from typing import Iterator
from typing import TypeVar
def prime_factors(n: int) -> Iterator[int]:
"""Return the prime factors of 'n' in order using a wheel with basis [2, 3, 5].
Translation of `primeFactors` from https://rosettacode.org/wiki/Category_talk:Wren-math.
"""
if n < 2:
return
inc = [4, 2, 4, 2, 4, 6, 2, 6]
while n % 2 == 0:
yield 2
n //= 2
while n % 3 == 0:
yield 3
n //= 3
while n % 5 == 0:
yield 5
n //= 5
k = 7
i = 0
while k * k <= n:
if n % k == 0:
yield k
n //= k
else:
k += inc[i]
i = (i + 1) % 8
if n > 1:
yield n
_T = TypeVar("_T")
def sliding_window(it: Iterable[_T], n: int) -> Iterator[tuple[_T, ...]]:
it = iter(it)
window = deque(islice(it, n - 1), maxlen=n)
for x in it:
window.append(x)
yield tuple(window)
def ruth_aaron_numbers_factors() -> Iterator[int]:
for (i, a), (_, b) in sliding_window(
((n, sum(prime_factors(n))) for n in count(2)), 2
):
if a == b:
yield i
def ruth_aaron_numbers_divisors() -> Iterator[int]:
for (i, a), (_, b) in sliding_window(
((n, sum(set(prime_factors(n)))) for n in count(2)), 2
):
if a == b:
yield i
def ruth_aaron_triples_factors() -> Iterator[int]:
for (i, a), (_, b), (_, c) in sliding_window(
((n, sum(prime_factors(n))) for n in count(2)), 3
):
if a == b == c:
yield i
def ruth_aaron_triples_divisors() -> Iterator[int]:
for (i, a), (_, b), (_, c) in sliding_window(
((n, sum(set(prime_factors(n)))) for n in count(2)), 3
):
if a == b == c:
yield i
if __name__ == "__main__":
print("First 30 Ruth-Aaron numbers (factors):")
pprint(list(islice(ruth_aaron_numbers_factors(), 30)), compact=True)
print("")
print("First 30 Ruth-Aaron numbers (divisors):")
pprint(list(islice(ruth_aaron_numbers_divisors(), 30)), compact=True)
print("")
print("First Ruth-Aaron triple (factors):")
print(next(ruth_aaron_triples_factors()))
print("")
print("First Ruth-Aaron triple (divisors):")
print(next(ruth_aaron_triples_divisors()))