CodeIQ No.2628
2つの正整数 $(a, b)$ が与えられたとき、
等式 $x^{2} + a^{2} = y^{2} + b^{2}$ を満たす正整数解 $(x, y)$ の全てについて、$x+y$ の総和を関数 $f(a, b)$ とします。
$f(a, b)$ を求めてください。
・$1 \le a \lt b \le 10^{5}$
$x^{2} + a^{2} = y^{2} + b^{2}$ となる $(x, y)$ を求めたい。
変数と定数に分離すると、
$x^{2} - y^{2} = b^{2} - a^{2}$
左辺を因数分解して、
$(x+y)(x-y) = b^{2} - a^{2}$
ここで $b^{2} - a^{2}$ を二正整数( $p,q$ とする。$p \lt q$ )の積で計算することができれば、
$\begin{aligned} \left\{ \begin{array}{l} x+y = q \\ x-y = p \end{array} \right. \end{aligned}$
なので、$x$ と $y$ を以下のように求めることができる。
$\begin{aligned} \left\{ \begin{array}{l} \displaystyle x = \frac{q+p}{2} \\ \\ \displaystyle y = \frac{q-p}{2} \end{array} \right. \end{aligned}$
ここで、$x$ や $y$ は $0$ であってはならない(つまり $p=q$ であってはならない)ことに注意すること。
(ΦωΦ)＜ソース
#!/usr/bin/python3
def divisors(n):
sq = int(n ** .5)
res = ((i, n//i) for i in range(1, sq+1) if n%i==0)
return res
def f(a, b):
res = 0
pq = b**2 - a**2
for p, q in divisors(pq):
if p == q:
continue
if (q + p) & 1 or (q - p) & 1:
continue
res += q # x+y=q
return res
###
from collections import namedtuple
import unittest
T = namedtuple('T', 'func, inntak, expected')
test_cases = [
T(f, (3, 10), 104),
T(f, (10, 50), 3500),
T(f, (20, 100), 15022),
T(f, (4, 5), 9),
###
T(f, (10, 26), 716),
T(f, (11, 389), 291886),
T(f, (123, 456), 294166),
T(f, (35672, 61243), 5082405064),
T(f, (71200, 82321), 2490724704),
T(f, (19126, 98765), 16043311668),
]
class TestSequence(unittest.TestCase):
pass
def test_generator(func, inntak, expected):
def _test(self):
self.assertEqual(func(*inntak), expected)
return _test
for case in test_cases:
name = 'test-{}-{}'.format(case.func.__name__, case.inntak)
test = test_generator(*case)
setattr(TestSequence, name, test)
suite = unittest.TestLoader().loadTestsFromTestCase(TestSequence)
testresult = unittest.TextTestRunner().run(suite)
assert testresult.wasSuccessful()
###
a, b = map(int, input().split())
res = f(a, b)
print(res)
---
$a$ や $b$ がもっと大きな値だったとする。
仮に $1 \le a \lt b \le 10^{10}$ だったとする。
この場合、上のソースの関数divisors()が遅いので、返事がなかなか返ってこなくなる。
そこで、divisors()をもっと効率的に実装する必要がある。
ひとつの案として、
「素因数分解した結果の素数の集合を全集合とし、それを2つの部分集合に分ける」
というのが考えられる。
このとき、素因数分解するのは $b^{2} - a^{2}$ であることに注目すると、
$b+a$ と $b-a$ のそれぞれを素因数分解し、それらの結果をマージすれば、$b^{2} - a^{2}$ を素因数分解するよりも計算量はぐっと小さくできる。
(ΦωΦ)＜ソース
#!/usr/bin/python3
# -*- coding: utf-8 -*-
# †
from collections import defaultdict
from functools import reduce
from itertools import product
from operator import mul
def factorize(n):
res = defaultdict(int)
d = 2
while d * d <= n:
while n % d == 0:
res[d] += 1
n //= d
d += 1
if n > 1:
res[n] += 1
return res
def merge(*dd):
res = defaultdict(int)
for d in dd:
for k, v in d.items():
res[k] += v
return res
def gen_divisors(a, b):
n = b*b - a*a
factors = merge(factorize(b+a), factorize(b-a))
keys = factors.keys() # 素数
vals = factors.values() # 冪乗(個数)
for p in product(*(range(v+1) for v in vals)):
res = reduce(mul, (x**y for x, y in zip(keys, p)))
if res < n//res: # 結局、関数fで等号時を弾くので、ここで弾いておく
yield res, n//res
def f(a, b):
res = 0
for p, q in gen_divisors(a, b):
if (q + p) & 1 or (q - p) & 1:
continue
res += q
return res
###
from collections import namedtuple
import unittest
T = namedtuple('T', 'func, inntak, expected')
test_cases = [
T(f, (3, 10), 104),
T(f, (10, 50), 3500),
T(f, (20, 100), 15022),
T(f, (4, 5), 9),
###
T(f, (10, 26), 716),
T(f, (11, 389), 291886),
T(f, (123, 456), 294166),
T(f, (35672, 61243), 5082405064),
T(f, (71200, 82321), 2490724704),
T(f, (19126, 98765), 16043311668),
###
T(f, (1, 10**10), 163688880842463429502),
T(f, (9, 10**10), 105925721330425759518),
T(f, (7, 9999999936), 99999998730000003990),
T(f, (12, 9999999955), 99999999110000001848),
T(f, (16, 9999999017), 133333307160001284080),
T(f, (11, 9998000141), 169344028062100814926),
T(f, (2, 9998000059), 157233566772634607620),
T(f, (2, 9998000147), 198016019985396859584),
T(f, (13, 9998000179), 163570929954618401814),
T(f, (16, 9998000263), 144820388526916915560),
T(f, (8, 9999400031), 149765240169422529784),
T(f, (12, 9999400031), 114272001199952686992),
T(f, (14, 9999400031), 171980104323475531012),
T(f, (11, 9998000057), 127427336246752144960),
T(f, (21, 9998000057), 81788730691229073300),
]
class TestSequence(unittest.TestCase):
pass
def test_generator(func, inntak, expected):
def _test(self):
self.assertEqual(func(*inntak), expected)
return _test
for case in test_cases:
name = 'test-{}-{}'.format(case.func.__name__, case.inntak)
test = test_generator(*case)
setattr(TestSequence, name, test)
suite = unittest.TestLoader().loadTestsFromTestCase(TestSequence)
testresult = unittest.TextTestRunner().run(suite)
assert testresult.wasSuccessful()
###
a, b = map(int, input().split())
res = f(a, b)
print(res)
(ΦωΦ)＜おしまい