from operator import add, mul
from math import sqrt, e

# Direct implementations of iterative improvement
def square_root(a):
    """Return the square root of a.

    >>> square_root(9)
    3.0
    """
    x = 1
    while x*x != a:
        x = square_root_update(x, a)
    return x

def square_root_update(x, a):
    return (x + a/x) / 2

def cube_root(a):
    """Return the cube root of a.

    >>> cube_root(27)
    3.0
    """
    x = 1
    while pow(x, 3) != a:
        x = cube_root_update(x, a)
    return x

def cube_root_update(x, a):
    return (2*x + a/(x*x)) / 3

# Iterative improvement
def golden_update(x):
    """The golden update rule returns a number closer to phi than
    x."""
    return 1/x + 1

def golden_test(x):
    """Return whether x is the golden ratio (phi)."""
    return x * x == x + 1

def iter_improve(update, done, guess=1, max_updates=1000):
    """Iteratively improve guess with update until done returns a true
    value.

    >>> iter_improve(golden_update, golden_test)
    1.618033988749895
    """
    k = 0
    while not done(guess) and k < max_updates:
        guess = update(guess)
        k = k + 1
    return guess

def square_root_improve(a):
    """Return the square root of a.

    >>> square_root_improve(9)
    3.0
    """
    def update(guess):
        return square_root_update(guess, a)
    def done(guess):
        return guess*guess == a
    return iter_improve(update, done)

def cube_root_improve(a):
    """Return the cube root of a.

    >>> cube_root_improve(27)
    3.0
    """
    def update(guess):
        return cube_root_update(guess, a)
    def done(guess):
        return pow(guess, 3) == a
    return iter_improve(update, done)

# Newton's method for nth roots
def nth_root_func_and_derivative(n, a):
    def root_func(x):
        return pow(x, n) - a
    def derivative(x):
        return n * pow(x, n-1)
    return root_func, derivative

def nth_root_newton(a, n):
    """Return the nth root of a.

    >>> nth_root_newton(8, 3)
    2.0
    >>> nth_root_newton(10, 4)
    1.7782794100389228
    """
    root_func, deriv = nth_root_func_and_derivative(n, a)
    def update(x):
        return x - root_func(x) / deriv(x)
    def done(x):
        return root_func(x) == 0
    return iter_improve(update, done)

# See online text for general Newton's method

# Recursion
def factorial_iter(n):
    """Compute the factorial of a natural number.

    >>> factorial(6)
    720
    """
    if n == 0 or n == 1:
        return 1
    total = 1
    while n >= 1:
        total, n = total * n, n - 1
    return total

def factorial(n):
    """Compute the factorial of a natural number.

    >>> factorial(6)
    720
    """
    if n == 0 or n == 1:
        return 1
    return n * factorial(n - 1)