def SegmentMax ( a , ids ) : func = lambda idxs : np . amax ( a [ idxs ] , axis = 0 ) return seg_map ( func , a , ids ) ,
idex of max in an array python
Segmented max op .
def SegmentMax(a, ids):
    """
    Segmented max op.
    """
    func = lambda idxs: np.amax(a[idxs], axis=0)
    return seg_map(func, a, ids),