高效分位箱

Question

我试图模仿pd.qcut的功能，但是它需要更有效。现在已基本实现，但无效值未正确处理。

import pandas as pd

def qcut(self,q,axis=-1):
    iaxs = np.ogrid[tuple(map(slice,self.shape))]
    cout = self.shape[axis]

    boud = np.ceil(np.linspace(0,cout,q+1)).astype(int)
    icut = np.empty(self.shape,dtype=uint16)
    inds = np.argpartition(self,boud[1:-1],axis=axis) 

    iaxs[axis] = inds
    icut[iaxs] = np.repeat(np.arange(q),np.diff(boud))
    return icut

aray = np.random.randn(10)

q1 = qcut(aray,3)
q2 = pd.qcut(aray,3,labels=False)
print q1
print q2
print (q1==q2).all()

输出：正确的输出

[0 1 1 0 2 0 1 2 0 2]
[0 1 1 0 2 0 1 2 0 2]
True

立即添加无效值：

aray[1] = np.nan
q1 = qcut(aray,3)*1.0; q1[1]=np.nan
q2 = pd.qcut(aray,3, labels=False)
print q1
print q2
print (q1==q2).all()

输出：输出不正确

[ 0.000  nan  1.000  0.000  2.000  0.000  1.000  2.000  0.000  1.000]
[ 0.000  nan  1.000  0.000  2.000  1.000  1.000  2.000  0.000  2.000]
False

我希望对解决方案的速度没有太大影响，谢谢