def initialize_counts(nS, nA):
    """Initializes a counts array.

    Parameters
    ----------
    nS: int
        Number of states
    nA: int
        Number of actions

    Returns
    -------
    counts: np.array of shape [nS x nA x nS]
        counts[state][action][next_state] is the number of times that doing "action" at state "state" transitioned to "next_state"
    """
    counts = [[[0 for _ in range(nS)] for _ in range(nA)] for _ in range(nS)]

    return counts