code: def density_to_peaks_vectorize(
      dists, nbrs, density, max_conn=1, name=""
  ):
      # just calculate 1 connectivity
      assert dists.shape[0] == density.shape[0]
      assert dists.shape == nbrs.shape
  
      num, k = dists.shape
  
      if name == "gcn_feat":
          include_mask = nbrs != np.arange(0, num).reshape(-1, 1)
          secondary_mask = (
              np.sum(include_mask, axis=1) == k
          )  # TODO: the condition == k should not happen as distance to the node self should be smallest, check for numerical stability; TODO: make top M instead of only supporting top 1
          include_mask[secondary_mask, -1] = False
          nbrs_exclude_self = nbrs[include_mask].reshape(
              -1, k - 1
          )  # (V, 79)
          dists_exclude_self = dists[include_mask].reshape(
              -1, k - 1
          )  # (V, 79)
      else:
          include_mask = nbrs != np.arange(0, num).reshape(-1, 1)
          nbrs_exclude_self = nbrs[include_mask].reshape(
              -1, k - 1
          )  # (V, 79)
          dists_exclude_self = dists[include_mask].reshape(
              -1, k - 1
          )  # (V, 79)
  
      compare_map = density[nbrs_exclude_self] > density.reshape(-1, 1)
      peak_index = np.argmax(
          np.where(compare_map, 1, 0), axis=1
      )  # (V,)
      compare_map_sum = np.sum(
          compare_map.cpu().data.numpy(), axis=1
      )  # (V,)
  
      dist2peak = {
          i: []
          if compare_map_sum[i] == 0
          else [dists_exclude_self[i, peak_index[i]]]
          for i in range(num)
      }
      peaks = {
          i: []
          if compare_map_sum[i] == 0
          else [nbrs_exclude_self[i, peak_index[i]]]
          for i in range(num)
      }
  
      return dist2peak, peaks
  
imports: from itertools import groupby
  import numpy
  import torch
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: density_to_peaks_vectorize(dists, nbrs, density, max_conn, name) and returns: dist2peak,  peaks