CodeJam: Cubic UFO

cubic_ufo.py

from math import sqrt, pow
import sys

def norm(v):
    m = mag(v)
    norm = tuple(n/m for n in v)
    return norm

def cross(a, b):
    c = [a[1]*b[2] - a[2]*b[1],
         a[2]*b[0] - a[0]*b[2],
         a[0]*b[1] - a[1]*b[0]]
    return c

def dot(a, b):
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]

def rotation_matrix(rcos, rsin, vec):
    u, v, w = vec
    matrix = [[0 for i in range(3)] for j in range(3)]
    matrix[0][0] =      rcos + u*u*(1-rcos)
    matrix[1][0] =  w * rsin + v*u*(1-rcos)
    matrix[2][0] = -v * rsin + w*u*(1-rcos)
    matrix[0][1] = -w * rsin + u*v*(1-rcos)
    matrix[1][1] =      rcos + v*v*(1-rcos)
    matrix[2][1] =  u * rsin + w*v*(1-rcos)
    matrix[0][2] =  v * rsin + u*w*(1-rcos)
    matrix[1][2] = -u * rsin + v*w*(1-rcos)
    matrix[2][2] =      rcos + w*w*(1-rcos)
    return matrix

def mag(v):
    return sqrt(sum(n*n for n in v))

f1 = [[0.5,0,0]]
f2 = [[0,0.5,0]]
f3 = [[0,0,0.5]]
corner = [0.5,0.5,0.5]
corner_unit = norm(corner)
testCases = int(input())
for t in range(testCases):
    A = float(input())
    if A == 1:
        print(" ".join(str(x) for x in f1[0]))
        print(" ".join(str(x) for x in f2[0]))
        print(" ".join(str(x) for x in f3[0]))
        continue
    Vdest = [sqrt(0.75 - pow(A/2, 2)), 0, A/2]
    Vdest_u = norm(Vdest)
    Vaxis = cross(corner_unit, Vdest_u)
    Vaxis_u = norm(Vaxis)
    rsin = mag(Vaxis) 
    rcos = dot(Vdest_u, corner_unit)
    rot = rotation_matrix(rcos, rsin, Vaxis_u)
    a1 = [[sum(row_el*col_el for row_el,col_el in zip(row,col)) for col in f1] for row in rot]
    a2 = [[sum(row_el*col_el for row_el,col_el in zip(row,col)) for col in f2] for row in rot]
    a3 = [[sum(row_el*col_el for row_el,col_el in zip(row,col)) for col in f3] for row in rot]
    print("Case #{}:".format(t+1))
    print(" ".join(str(x[0]) for x in a1))
    print(" ".join(str(x[0]) for x in a2))
    print(" ".join(str(x[0]) for x in a3))