nodes: 0, module; 1, function_definition; 2, function_name:argmaxMulti; 3, parameters; 4, block; 5, identifier:a; 6, identifier:groupKeys; 7, default_parameter; 8, expression_statement; 9, if_statement; 10, expression_statement; 11, expression_statement; 12, expression_statement; 13, comment:# Break ties by finding the insertion points of the the group start indices; 14, comment:# and using the values currently at those points. This approach will choose; 15, comment:# the first occurrence of each max value.; 16, expression_statement; 17, if_statement; 18, identifier:assumeSorted; 19, False; 20, comment:"""
    This is like numpy's argmax, but it returns multiple maximums.
  
    It gets the indices of the max values of each group in 'a', grouping the
    elements by their corresponding value in 'groupKeys'.
  
    @param a (numpy array)
    An array of values that will be compared
  
    @param groupKeys (numpy array)
    An array with the same length of 'a'. Each entry identifies the group for
    each 'a' value.
  
    @param assumeSorted (bool)
    If true, group keys must be organized together (e.g. sorted).
  
    @return (numpy array)
    The indices of one maximum value per group
  
    @example
      _argmaxMulti([5, 4, 7, 2, 9, 8],
                   [0, 0, 0, 1, 1, 1])
    returns
      [2, 4]
    """; 21, not_operator; 22, comment:# Use a stable sort algorithm; 23, block; 24, assignment; 25, assignment; 26, assignment; 27, assignment; 28, identifier:assumeSorted; 29, block; 30, else_clause; 31, identifier:assumeSorted; 32, expression_statement; 33, expression_statement; 34, expression_statement; 35, pattern_list; 36, call; 37, identifier:maxValues; 38, call; 39, identifier:allMaxIndices; 40, call; 41, identifier:indices; 42, subscript; 43, return_statement; 44, block; 45, assignment; 46, assignment; 47, assignment; 48, identifier:_; 49, identifier:indices; 50, identifier:lengths; 51, attribute; 52, argument_list; 53, attribute; 54, argument_list; 55, attribute; 56, argument_list; 57, identifier:allMaxIndices; 58, call; 59, identifier:indices; 60, return_statement; 61, identifier:sorter; 62, call; 63, identifier:a; 64, subscript; 65, identifier:groupKeys; 66, subscript; 67, identifier:np; 68, identifier:unique; 69, identifier:groupKeys; 70, keyword_argument; 71, keyword_argument; 72, attribute; 73, identifier:reduceat; 74, identifier:a; 75, identifier:indices; 76, identifier:np; 77, identifier:flatnonzero; 78, comparison_operator:np.repeat(maxValues, lengths) == a; 79, attribute; 80, argument_list; 81, subscript; 82, attribute; 83, argument_list; 84, identifier:a; 85, identifier:sorter; 86, identifier:groupKeys; 87, identifier:sorter; 88, identifier:return_index; 89, True; 90, identifier:return_counts; 91, True; 92, identifier:np; 93, identifier:maximum; 94, call; 95, identifier:a; 96, identifier:np; 97, identifier:searchsorted; 98, identifier:allMaxIndices; 99, identifier:indices; 100, identifier:sorter; 101, identifier:indices; 102, identifier:np; 103, identifier:argsort; 104, identifier:groupKeys; 105, keyword_argument; 106, attribute; 107, argument_list; 108, identifier:kind; 109, string:"mergesort"; 110, identifier:np; 111, identifier:repeat; 112, identifier:maxValues; 113, identifier:lengths
code: def argmaxMulti(a, groupKeys, assumeSorted=False):
    """
    This is like numpy's argmax, but it returns multiple maximums.
  
    It gets the indices of the max values of each group in 'a', grouping the
    elements by their corresponding value in 'groupKeys'.
  
    @param a (numpy array)
    An array of values that will be compared
  
    @param groupKeys (numpy array)
    An array with the same length of 'a'. Each entry identifies the group for
    each 'a' value.
  
    @param assumeSorted (bool)
    If true, group keys must be organized together (e.g. sorted).
  
    @return (numpy array)
    The indices of one maximum value per group
  
    @example
      _argmaxMulti([5, 4, 7, 2, 9, 8],
                   [0, 0, 0, 1, 1, 1])
    returns
      [2, 4]
    """
  
    if not assumeSorted:
      # Use a stable sort algorithm
      sorter = np.argsort(groupKeys, kind="mergesort")
      a = a[sorter]
      groupKeys = groupKeys[sorter]
  
    _, indices, lengths = np.unique(groupKeys, return_index=True,
                                    return_counts=True)
  
    maxValues = np.maximum.reduceat(a, indices)
    allMaxIndices = np.flatnonzero(np.repeat(maxValues, lengths) == a)
  
    # Break ties by finding the insertion points of the the group start indices
    # and using the values currently at those points. This approach will choose
    # the first occurrence of each max value.
    indices = allMaxIndices[np.searchsorted(allMaxIndices, indices)]
  
    if assumeSorted:
      return indices
    else:
      return sorter[indices]