code: def _stratified_split(
      interactions: BaseInteractions,
      test_p: float,
      processes: int,
      seed: int,
      force_split: bool,
  ) -> Tuple[Interactions, Interactions]:
      users = interactions.mat.row
      unique_users = set(users)
  
      # while we should be able to run ``np.where(users == user)[0]`` to find all items each user
      # interacted with, by building up a dictionary to get these values instead, we can achieve the
      # same result in O(N) complexity rather than O(M * N), a nice timesave to have when working with
      # larger datasets
      all_idxs_for_users_dict = defaultdict(list)
      for idx, user in enumerate(users):
          all_idxs_for_users_dict[user].append(idx)
  
      if processes == 0:
          test_idxs = [
              _stratified_split_parallel_worker(
                  idxs_to_split=all_idxs_for_users_dict[user],
                  test_p=test_p,
                  seed=(seed + user),
                  force_split=force_split,
              )
              for user in unique_users
          ]
      else:
          # run the function below in parallel for each user
          # by setting the seed to ``seed + user``, we get a balance between reproducability and
          # actual randomness so users with the same number of interactions are not split the exact
          # same way
          test_idxs = Parallel(n_jobs=processes)(
              delayed(_stratified_split_parallel_worker)(
                  all_idxs_for_users_dict[user],
                  test_p,
                  seed + user,
                  force_split,
              )
              for user in unique_users
          )
  
      # reduce the list of lists down to a 1-d list
      test_idxs = functools.reduce(operator.iconcat, test_idxs, [])
      # find all indices not in test set - they are now train
      train_idxs = list(set(range(len(users))) - set(test_idxs))
  
      train_interactions = _subset_interactions(
          interactions=interactions, idxs=train_idxs
      )
      test_interactions = _subset_interactions(
          interactions=interactions, idxs=test_idxs
      )
  
      return train_interactions, test_interactions
  
file_path: collie/collie/cross_validation.py
imports: from collections import defaultdict
  import functools
  import operator
  from typing import Any, Iterable, Optional, Tuple, Union
  from joblib import delayed, Parallel
  import numpy
  from scipy.sparse import coo_matrix
  from sklearn.model_selection import train_test_split
  from collie.interactions import BaseInteractions, ExplicitInteractions, HDF5Interactions, Interactions
  from collie.utils import get_random_seed
label_desc: Use image node_id 4 for calling a global function with example usage: _stratified_split(interactions, test_p, processes, seed, force_split) and returns: train_interactions,  test_interactions