multiprocess brute force 596

pe_596_2.py

#!/usr/bin/env python3
"""
Project Euler, Problem 596
https://projecteuler.net/problem=596

Let be T(r) the number of integer quadruplets x, y, z, t such that x**2 + y**2 + z**2 + t**2 ≤ r**2.
In other words, T(r) is the number of lattice points in the four-dimensional hyperball of radius r.

You are given that T(2) = 89, T(5) = 3121, T(100) = 493490641 and T(10**4) = 49348022079085897.

Find T(10**8) mod 1000000007.
"""

from multiprocessing import Pool, cpu_count
from functools import partial


def calc_inner_4cube(in_r):
    return in_r**4  # (2*in_r//2)**4


def x_loop(x_axis, in_r):
    ring = in_r**2
    result = 0
    temp_ring_y = int((ring - x_axis**2)**.5)
    for y_axis in range(x_axis, temp_ring_y + 1):
        temp_ring_z = int((temp_ring_y - y_axis**2)**.5)
        for z_axis in range(y_axis, temp_ring_z + 1):
            temp_ring_t = int((temp_ring_z - z_axis**2)**.5)
            for t_axis in range(z_axis, temp_ring_t + 1):
                if ring >= x_axis**2 + y_axis**2 + z_axis**2 + t_axis**2:
                    result += 1
    return result
                        
    
def mp_brute_calc(in_r):
    """iterate through hypercube of radius r"""
    ring = in_r**2
    result = calc_inner_4cube(in_r)
    worker_x = partial(x_loop, in_r=in_r)
    with Pool(cpu_count() - 1) as poolio:
        result += sum(poolio.map(worker_x, range(in_r // 2, in_r + 1)))
    return result


if __name__ == '__main__':
    assert mp_brute_calc(2) == 89
    assert mp_brute_calc(5) == 3121
    assert mp_brute_calc(100) == 493490641
    assert mp_brute_calc(104) == 49348022079085897