code: def twostep_count_matrix(dtrajs, lag, N):
      
      # List all transition triples:
      rows = []
      cols = []
      states = []
      for dtraj in dtrajs:
          if dtraj.size > 2*lag:
              rows.append(dtraj[0:-2*lag])
              states.append(dtraj[lag:-lag])
              cols.append(dtraj[2*lag:])
      row = np.concatenate(rows)
      col = np.concatenate(cols)
      state = np.concatenate(states)
      data = np.ones(row.size)
      # Transform the rows and cols into a single list with N*+2 possible values:
      pair = N * row + col
      # Estimate sparse matrix:
      C2t = scipy.sparse.coo_matrix((data, (pair, state)), shape=(N*N, N))
  
      return C2t.tocsc()
docs: Compute all two-step count matrices from discrete trajectories.
  
      Parameters
      ----------
      dtrajs : list of discrete trajectories
  
      lag : int
          the lag time for count matrix estimation
      N : int
          the number of states in the discrete trajectories.
  
      Returns
      -------
      C2t : sparse csc-matrix (N, N, N)
          two-step count matrices for all states. C2t[:, n, :] is a count matrix for each n