python - 有没有办法将 numpy 数组映射到某个数据框？

Question

我有一个 DataFrame，比方说:

#d = {'col1': [1, 2, 3], 'col2': [3, 4, 5]} // that's what the data might look like
df = pd.DataFrame(data=d)

我有一个带有 [0, 2] 的 np 数组。

现在我想向 DataFrame 添加一列，当行的索引在 np 数组中时为 1，否则为 0。

有人有想法吗？

Answer 1

使用Index.isin将掩码转换为整数:

d = {'col1': [1, 2, 3], 'col2': [3, 4, 5]} 
df = pd.DataFrame(data=d)

a = np.array([0, 2])

df['new'] = df.index.isin(a).astype(int)
#alternative
#df['new'] = np.in1d(df.index, a).astype(int)

或者使用numpy.where :

df['new'] = np.where(df.index.isin(a), 1, 0)
#alternative
#df['new'] = np.where(np.in1d(df.index, a), 1, 0)

print (df)
   col1  col2  new
0     1     3    1
1     2     4    0
2     3     5    1