class_docstr: _EfficientKMeans is primarily used by palettization to perform a k-means operation. This class also has an in-house
  implementation of k-means, called ``kmeans_pp`` which runs entirely on GPU and is ~10x faster than sklearn's API.
docstr: Search for row-maxima in a 2d totally monotone matrix M[i,j].
  The input is specified by a list of row indices, a list of column
  indices, and a function "lookup" satisfying lookup(i,j) = M[i,j].
  The matrix must satisfy the totally monotone ordering property:
  if i occurs before i' in rows, j occurs before j' in cols, and
  M[i,j] < M[i,j'], then also M[i',j] < M[i',j'].  The result is
  returned as a dictionary mapping row i to the column j containing
  the largest value M[i,j].  Ties are broken in favor of earlier
  columns.  The number of calls to lookup is O(len(rows)+len(cols)).
docstr_tok: ["Search","for","row-maxima","in","a","2d","totally","monotone","matrix","M","[","i",",","j","]",".","The","input","is","specified","by","a","list","of","row","indices",",","a","list","of","column","indices",",","and","a","function","``","lookup","''","satisfying","lookup","(","i",",","j",")","=","M","[","i",",","j","]",".","The","matrix","must","satisfy","the","totally","monotone","ordering","property",":","if","i","occurs","before","i","'","in","rows",",","j","occurs","before","j","'","in","cols",",","and","M","[","i",",","j","]","<","M","[","i",",","j","'","]",",","then","also","M","[","i","'",",","j","]","<","M","[","i","'",",","j","'","]",".","The","result","is","returned","as","a","dictionary","mapping","row","i","to","the","column","j","containing","the","largest","value","M","[","i",",","j","]",".","Ties","are","broken","in","favor","of","earlier","columns",".","The","number","of","calls","to","lookup","is","O","(","len","(","rows",")","+len","(","cols",")",")","."]
code: def smawk(self, rows, cols, lookup):
      """Search for row-maxima in a 2d totally monotone matrix M[i,j].
      The input is specified by a list of row indices, a list of column
      indices, and a function "lookup" satisfying lookup(i,j) = M[i,j].
      The matrix must satisfy the totally monotone ordering property:
      if i occurs before i' in rows, j occurs before j' in cols, and
      M[i,j] < M[i,j'], then also M[i',j] < M[i',j'].  The result is
      returned as a dictionary mapping row i to the column j containing
      the largest value M[i,j].  Ties are broken in favor of earlier
      columns.  The number of calls to lookup is O(len(rows)+len(cols)).
      """
  
      # base case of recursion
      if not rows:
          return {}
  
      # reduce phase: make number of columns at most equal to number of rows
      stack = []
      for c in cols:
          while len(stack) >= 1 and lookup(
              rows[len(stack) - 1], stack[-1]
          ) < lookup(rows[len(stack) - 1], c):
              stack.pop()
          if len(stack) != len(rows):
              stack.append(c)
  
      cols = stack
  
      # recursive call to search for every odd row
      result = self.smawk(
          [rows[i] for i in range(1, len(rows), 2)], cols, lookup
      )
  
      # go back and fill in the even rows
      c = 0
      for r in range(0, len(rows), 2):
          row = rows[r]
          if r == len(rows) - 1:
              cc = (
                  len(cols) - 1
              )  # if r is last row, search through last col
          else:
              cc = c  # otherwise only until pos of max in row r+1
              target = result[rows[r + 1]]
              while cols[cc] != target:
                  cc += 1
          result[row] = max(
              [
                  (lookup(row, cols[x]), -x, cols[x])
                  for x in range(c, cc + 1)
              ]
          )[2]
          c = cc
  
      return result
  
file_path: coremltools/coremltools/optimize/torch/palettization/_efficient_kmeans.py
imports: import numpy
  import torch
  import torch.distributed
label_desc: Use image node_id 10 for calling the _EfficientKMeans obj's underlying member method code with example usage: obj.smawk(rows, cols, lookup) and returns: result, dict