python - Pandas 根据索引值分配标签

Question

我有一个带有索引和多列的数据框。其次，我有几个包含根据某些标准采样的索引值的列表。现在我想根据特定行的索引是否存在于指定列表中的事实来创建带有标签的列。

现在我使用它有两种情况:

1) 创建一列并根据一个列表给出标签:

df['1_name'] = df.index.map(lambda ix: 'A' if ix in idx_1_model else 'B')

2) 创建列并根据多个列表给出标签:

def assignLabelsToSplit(ix_, random_m, random_y, model_m, model_y):
    if (ix_ in random_m) or (ix_ in model_m):
        return 'A'
    if (ix_ in random_y) or (ix_ in model_y):
        return 'B'
    else:
        return 'not_assigned'

df['2_name'] = df.index.map(lambda ix: assignLabelsToSplit(ix, idx_2_random_m, idx_2_random_y, idx_2_model_m, idx_2_model_y))

这可以工作，但速度很慢。每次调用大约需要 3 分钟，考虑到我必须多次执行函数，它需要更快。

感谢您的任何建议。

Answer 1

我认为你需要双 numpy.where与 Index.isin :

df['2_name'] = np.where(df.index.isin(random_m + model_m), 'A',
               np.where(df.index.isin(random_y + model_y), 'B', 'not_assigned'))

示例:

np.random.seed(100)
df = pd.DataFrame(np.random.randint(10, size=(10,1)), columns=['A'])
#print (df)

random_m = [0,1]
random_y =  [2,3]
model_m = [7,4]
model_y = [5,6]

print (type(random_m))
<class 'list'>

print (random_m + model_m)
[0, 1, 7, 4]

print (random_y + model_y)
[2, 3, 5, 6]

df['2_name'] = np.where(df.index.isin(random_m + model_m), 'A',
               np.where(df.index.isin(random_y + model_y), 'B', 'not_assigned'))
print (df)
   A        2_name
0  8             A
1  8             A
2  3             B
3  7             B
4  7             A
5  0             B
6  4             B
7  2             A
8  5  not_assigned
9  2  not_assigned