docs: def iterative_stratification(node_label_matrix, training_set_size, number_of_categories, random_seed=0):
      """
      Iterative data fold stratification/balancing for two folds.
  
      Based on: Sechidis, K., Tsoumakas, G., & Vlahavas, I. (2011).
                On the stratification of multi-label data.
                In Machine Learning and Knowledge Discovery in Databases (pp. 145-158).
                Springer Berlin Heidelberg.
  
      Inputs:  - node_label_matrix: The node-label ground truth in a SciPy sparse matrix format.
               - training_set_size: The minimum required size for the training set.
               - number_of_categories: The number of categories/classes in the learning.
               - random_seed: A seed for numpy random.
  
      Outputs: - train_set: A NumPy array containing the training set node ids.
               - test_set: A NumPy array containing the testing set node ids.
      """
      number_of_labelled_nodes = node_label_matrix.shape[0]
      testing_set_size = number_of_labelled_nodes - training_set_size
      training_set_proportion = training_set_size/number_of_labelled_nodes
      testing_set_proportion = testing_set_size/number_of_labelled_nodes
  
      # Calculate the desired number of examples of each label at each subset.
      desired_label_number = np.zeros((2, number_of_categories), dtype=np.int64)
      node_label_matrix = node_label_matrix.tocsc()
      for j in range(number_of_categories):
          category_label_number = node_label_matrix.getcol(j).indices.size
          desired_label_number[0, j] = math.ceil(category_label_number*training_set_proportion)
          desired_label_number[1, j] = category_label_number - desired_label_number[0, j]
  
      train_ids = list()
      test_ids = list()
  
      append_train_id = train_ids.append
      append_test_id = test_ids.append
  
      # Randomize process
      np.random.seed(random_seed)
      while True:
          if len(train_ids) + len(test_ids) >= number_of_labelled_nodes:
              break
          # Find the label with the fewest (but at least one) remaining examples, breaking the ties randomly
          remaining_label_distribution = desired_label_number.sum(axis=0)
          min_label = np.min(remaining_label_distribution[np.where(remaining_label_distribution > 0)[0]])
          label_indices = np.where(remaining_label_distribution == min_label)[0]
          chosen_label = int(np.random.choice(label_indices, 1)[0])
  
          # Find the subset with the largest number of desired examples for this label,
          # breaking ties by considering the largest number of desired examples, breaking further ties randomly.
          fold_max_remaining_labels = np.max(desired_label_number[:, chosen_label])
          fold_indices = np.where(desired_label_number[:, chosen_label] == fold_max_remaining_labels)[0]
          chosen_fold = int(np.random.choice(fold_indices, 1)[0])
  
          # Choose a random example for the selected label.
          relevant_nodes = node_label_matrix.getcol(chosen_label).indices
          chosen_node = int(np.random.choice(np.setdiff1d(relevant_nodes,
                                                          np.union1d(np.array(train_ids),
                                                                     np.array(test_ids))),
                                             1)[0])
          if chosen_fold == 0:
              append_train_id(chosen_node)
              desired_label_number[0, node_label_matrix.getrow(chosen_node).indices] -= 1
          elif chosen_fold == 1:
              append_test_id(chosen_node)
              desired_label_number[1, node_label_matrix.getrow(chosen_node).indices] -= 1
          else:
              raise RuntimeError
  
      return np.array(train_ids), np.array(test_ids)