<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def number_of_states(dtrajs, only_used = False):
 r"""returns the number of states in the given trajectories. Parameters dtraj : array_like or list of array_like Discretized trajectory or list of discretized trajectories only_used = False : boolean If False, will return max+1, where max is the largest index used. If True, will return the number of states that occur at least once. """

    dtrajs = _ensure_dtraj_list(dtrajs)
    if only_used:
        # only states with counts > 0 wanted. Make a bincount and count nonzeros
        bc = count_states(dtrajs)
        return np.count_nonzero(bc)
    else:
        # all states wanted, included nonpopulated ones. return max + 1
        imax = 0
        for dtraj in dtrajs:
            imax = max(imax, np.max(dtraj))
        return imax+1