def argmaxMulti(a, groupKeys, assumeSorted=False):
  if not assumeSorted:
    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)
  indices = allMaxIndices[np.searchsorted(allMaxIndices, indices)]
  if assumeSorted:
    return indices
  else:
    return sorter[indices]