docs: def index_states(dtrajs, subset = None):
      """Generates a trajectory/time indexes for the given list of states
  
      Parameters
      ----------
      dtraj : array_like or list of array_like
          Discretized trajectory or list of discretized trajectories
      subset : ndarray((n)), optional, default = None
          array of states to be indexed. By default all states in dtrajs will be used
  
      Returns
      -------
      indexes : list of ndarray( (N_i, 2) )
          For each state, all trajectory and time indexes where this state occurs.
          Each matrix has a number of rows equal to the number of occurances of the corresponding state,
          with rows consisting of a tuple (i, t), where i is the index of the trajectory and t is the time index
          within the trajectory.
  
      """
      # check input
      dtrajs = _ensure_dtraj_list(dtrajs)
      # select subset unless given
      n = number_of_states(dtrajs)
      if subset is None:
          subset = list(range(n))
      else:
          if np.max(subset) >= n:
              raise ValueError('Selected subset is not a subset of the states in dtrajs.')
      # histogram states
      hist = count_states(dtrajs)
      # efficient access to which state are accessible
      is_requested = np.ndarray((n), dtype=bool)
      is_requested[:] = False
      is_requested[subset] = True
      # efficient access to requested state indexes
      full2states = np.zeros((n), dtype=int)
      full2states[subset] = list(range(len(subset)))
      # initialize results
      res    = np.ndarray((len(subset)), dtype=object)
      counts = np.zeros((len(subset)), dtype=int)
      for i,s in enumerate(subset):
          res[i] = np.zeros((hist[s],2), dtype=int)
      # walk through trajectories and remember requested state indexes
      for i,dtraj in enumerate(dtrajs):
          for t,s in enumerate(dtraj):
              if is_requested[s]:
                  k = full2states[s]
                  res[k][counts[k],0] = i
                  res[k][counts[k],1] = t
                  counts[k] += 1
      return res