Codejam 2018 Qualification Round Problem 4a

cube.py

""" We needed to dust the math textbooks! """
from math import sin
from math import cos
from math import sqrt
from math import pi
from math import atan

def reader():
    """Read the input (desired cube area)"""
    N = int(input())
    lst = []
    for i in range(N):
        lst.append(float(input()))
    return lst

def area(a, b):
    """It seems that the area is just the following formulae"""
    return sin(a) + sqrt(2)*cos(a)*cos(b-pi/4)

def solver(A):
    """ Solves area(a, b) for Area < sqrt(2) """
    b = 0 if A < sqrt(2) else pi/4
    low, high, mid = 0, atan(sqrt(2)/2) if b > 0 else pi/4, pi/8

    cnt = 50
    
    while abs(A - area(mid, b)) > 10e-8 and not cnt == 0:
        """ Binary search in the angles! 
        area(a, b) is ascending for both a, b in (0, sqrt(2)/2), (0, pi/4)
        """
        if area(mid, b) > A:
            high = mid
            mid = (mid + low)/2
        else:
            low = mid
            mid = (high + mid)/2
        cnt -=1
    # Wrong 2nd rotation (needed by x) missed me 21 points!
    return (rotate_by_z(rotate_by_z((.5, 0, 0), mid), b),
            rotate_by_z(rotate_by_z((0, .5, 0), mid), b),
            rotate_by_z(rotate_by_z((0, 0, .5), mid), b))

def rotate_by_z(P, theta):
    x, y, z = P
    return x*cos(theta) - y*sin(theta), x*sin(theta) + y*cos(theta), z

def rotate_by_y(P, theta):
    x, y, z = P
    return x*cos(theta) + z*sin(theta), y, -x*sin(theta) + z*cos(theta)

def rotate_by_x(P, theta):
    x, y, z = P
    return x, -z*sin(theta) + y*cos(theta), z*cos(theta) + y*sin(theta)


import sys
for i, L in enumerate(reader()):
    print("Case #{}:".format(i+1))
    for i in solver(L):
        print("{} {} {}".format(*i))
    sys.stdout.flush()