def identity(n):
    """
    Return the identity matrix of size n
    :param n: size of the identity matrix to return
    :return: identity matrix represented by a 2 dimensional array of size n*n
    """
    return [[1 if i == j else 0 for i in range(n)] for j in range(n)]