code: def calculate_removable_counts(
      removable_counts_mat: torch.Tensor,
      remain_count: int,
      num_clus: int,
  ) -> torch.Tensor:
      """
      Calculate removable counts based on the arguments and calculate how many counts should be
      removed from the each cluster. This function has `O(N)` (N = num_clus) time complexity to
      return the desired `removable_counts_mat`.
  
      Example:
  
          The original input to `get_merge_quantity` function:
          >>> pre_clus_labels = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2]
          >>> num_to_be_removed = 3
          >>> min_count_per_cluster = 2
  
          Histogram: (`min_count_per_cluster`=2 is removed)
          0 |*****
          1 |***
          2 |*
  
          Inputs:
              >>> removable_counts_mat = [5, 3, 1]
              >>> remain_count = 6
              >>> num_clus = 3
  
          Interim results:
              >>> diff_counts
              [1, 2, 2]
              >>> gradual_counts
              [3, 4, 2]
              >>> cumsum_counts
              [3, 7, 9]
  
          Return:
              >>> removable_counts_mat
              [2, 1, 0]
  
      Args:
          removable_counts_mat (Tensor):
              Tensor containing how many vectors could be removed from each cluster
          remain_count (int):
              Integer value that indicates the number of vectors removed from the total set
          num_clus (int):
              Number of clusters in the given label sequence (cardinality of a label set)
  
      Returns:
          removable_counts_mat (Tensor):
              Tensor containing the number of vectors should be removed from each cluster
      """
      device = removable_counts_mat.device
      zero_padded_counts = torch.cat(
          [
              torch.tensor([0]).to(device),
              removable_counts_mat.sort()[0],
              torch.tensor([0]).to(device),
          ],
          dim=0,
      )
      removable_count_args = removable_counts_mat.sort(descending=True)[
          1
      ]
  
      # Calculate the size difference between clusters
      diff_counts = (zero_padded_counts[1:] - zero_padded_counts[:-1])[
          :num_clus
      ]
      gradual_counts = (
          torch.arange(num_clus, 0, -1).to(device) * diff_counts
      )
      cumsum_counts = torch.cumsum(gradual_counts, dim=0)
      remain_count_rem = remain_count
  
      # Find how many remaining counts we can use
      ind: int = 0
      for ind, num in enumerate(cumsum_counts):
          if remain_count < num:
              break
  
      # Subtract the common values step by step
      if ind > 0:
          for knd in range(ind):
              removable_counts_mat[
                  removable_count_args[: num_clus - knd]
              ] -= diff_counts[knd]
              remain_count_rem -= int(diff_counts[knd].item()) * (
                  num_clus - knd
              )
      assert remain_count >= 0, "remain_count should never be negative."
  
      # Add remaining values
      num_labels = remain_count_rem // (num_clus - ind)
      rem_labels = remain_count_rem % (num_clus - ind)
      removable_counts_mat[
          removable_count_args[: (num_clus - ind)]
      ] -= num_labels
      removable_counts_mat[removable_count_args[:rem_labels]] -= 1
      return removable_counts_mat.int()
  
imports: from typing import List, Tuple
  import torch
  from nemo.collections.asr.parts.utils.offline_clustering import NMESC, SpeakerClustering, SpectralClustering, get_scale_interpolated_embs, getAffinityGraphMat, getCosAffinityMatrix, split_input_data
  from nemo.collections.asr.parts.utils.optimization_utils import linear_sum_assignment
label_desc: Use image node_id 4 for calling a global function with example usage: calculate_removable_counts(removable_counts_mat, remain_count, num_clus) and returns: removable_counts_mat