如何有效地将值列表连接到间隔列表?

时间:2014-02-26 01:50:17

标签: python numpy pandas

我有一个数据框,可以按如下方式构建:

df = pd.DataFrame({'value':scipy.stats.norm.rvs(0, 1, size=1000), 
           'start':np.abs(scipy.stats.norm.rvs(0, 20, size=1000))})
df['end'] = df['start'] + np.abs(scipy.stats.norm.rvs(5, 5, size=1000))
df[:10]

       start       value        end
0    9.521781   -0.570097    17.708335
1    3.929711   -0.927318    15.065047
2    3.990466    0.756413    4.841934
3    20.676291  -1.418172    28.284301
4    13.084246   1.280723    14.121626
5    29.784740   0.236915    32.791751
6    21.626625   1.144663    28.739413
7    18.524309   0.101871    27.271344
8    21.288152  -0.727120    27.049582
9    13.556664   0.713141    22.136275

每行代表一个分配给间隔(开始,结束)的值

现在,我想得到一个在10,13,15,...,70时出现的最佳值列表。 (如果您熟悉它,它类似于SQL中的几何索引。)

以下是我第一次尝试使用pandas的python,需要18.5ms。任何人都可以帮助改善吗? (此程序将在我的程序中使用不同的数据框调用1M或更多次)

def get_values(data):
    data.sort_index(by='value', ascending=False, inplace=True) # this takes 0.2ms
    # can we get rid of it? since we don't really need sort...
    # all we need is the max value for each interval. 
    # But if we have to keep it for simplicity it is ok.
    ret = []
    #data = data[(data['end'] >= 10) & (data['start'] <= 71)]
    for t in range(10, 71, 2):
        interval = data[(data['end'] >= t) & (data['start'] <= t)]
        if not interval.empty:
            ret.append(interval['value'].values[0])
        else:
            for i in range(t, 71, 2):
                ret.append(None)
            break
    return ret
#%prun -l 10 print get_values(df)
%timeit get_values(df)

第二次尝试涉及尽可能将大熊猫分解成numpy,需要大约0.7ms

def get_values(data):
    data.sort_index(by='value', ascending=False, inplace=True)
    ret = []
    df_end = data['end'].values
    df_start = data['start'].values
    df_value = data['value'].values
    for t in range(10, 71, 2):
        values = df_value[(df_end >= t) & (df_start <= t)]
        if len(values) != 0:
            ret.append(values[0])
        else:
            for i in range(t, 71, 2):
                ret.append(None)
            break
    return ret
#%prun -l 10 print get_values(df)
%timeit get_values(df)

我们可以进一步提高吗?我想下一步是算法级别,上面两个都只是天真的逻辑实现。

1 个答案:

答案 0 :(得分:2)

我不理解你的代码中的空进程,如果忽略你的空进程,这里是一个更快的版本:

import scipy.stats as stats
import pandas as pd
import numpy as np

df = pd.DataFrame({'value':stats.norm.rvs(0, 1, size=1000), 
           'start':np.abs(stats.norm.rvs(0, 20, size=1000))})
df['end'] = df['start'] + np.abs(stats.norm.rvs(5, 5, size=1000))

def get_value(df, target):
    value = df["value"].values
    idx = np.argsort(value)[::-1]
    start = df["start"].values[idx]
    end = df["end"].values[idx]
    value = value[idx]

    mask = (target[:, None] >= start[None, :]) & (target[:, None] <= end[None, :])
    index = np.argmax(mask, axis=1)
    flags = mask[np.arange(len(target)), index]
    result = value[index]
    result[~flags] = np.nan
    return result

get_value(df, np.arange(10, 71, 2))