Vectorizing a Numpy slice operation

后端 未结 3 990
梦如初夏
梦如初夏 2021-01-13 00:53

Say I have a Numpy vector,

A = zeros(100)

and I divide it into subvectors by a list of breakpoints which index into A, for ins

3条回答
  •  春和景丽
    2021-01-13 01:22

    You can use simple np.cumsum -

    import numpy as np
    
    # Form zeros array of same size as input array and 
    # place ones at positions where intervals change
    A1 = np.zeros_like(A)
    A1[breaks[1:-1]] = 1
    
    # Perform cumsum along it to create a staircase like array, as the final output
    out = A1.cumsum()
    

    Sample run -

    In [115]: A
    Out[115]: array([3, 8, 0, 4, 6, 4, 8, 0, 2, 7, 4, 9, 3, 7, 3, 8, 6, 7, 1, 6])
    
    In [116]: breaks
    Out[116]: array([ 0,  4,  9, 11, 18, 20])
    
    In [142]: out
    Out[142]: array([0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4]..)
    

    If you want to have mean values of those subvectors from A, you can use np.bincount -

    mean_vals = np.bincount(out, weights=A)/np.bincount(out)
    

    If you are looking to extend this functionality and use a custom function instead, you might want to look into MATLAB's accumarray equivalent for Python/Numpy: accum whose source code is available here.

提交回复
热议问题