def get_policy(Q):
    """
    Computes an optimal policy from a Q-matrix.
    Args:
        - Q(array): nS*nA Q-matrix.
    Returns:
        a policy matrix (nS*nA), assigning uniform probability across all
        optimal actions in a given state.
    """
    nS, nA = Q.shape
    pi = (Q >= Q.max(1).reshape(nS, 1))
    # If more than one optimal action, split probability between them
    pi = pi / pi.sum(1).reshape(nS, 1)
    return pi