用包含数组条目的列查询熊猫

时间:2018-07-23 23:59:28

标签: python pandas dataframe series

ykp.data
Out[182]: 
    state  action  reward  
0    [41]       5      59  
1     [5]      52      48  
2    [46]      35      59  
3    [42]      16      12  
4    [43]      37      48   
5    [36]       5      59   
6    [49]      52      48 
7    [39]      11      23 

我想在状态条目中找到与[42]匹配的行,所以我跑了

ykp.data.query('state == [42]')

但我知道

Empty DataFrame
Columns: [state, action, reward]
Index: []

当我应该看到[42], 16, 12时。

有人可以告诉我如何解决此问题吗?我需要将状态值存储为数组。

3 个答案:

答案 0 :(得分:4)

在此最好避免使用pd.Series.apply。相反,您可以使用itertools.chain来构造一个常规的NumPy数组。然后将数组与整数进行比较,以形成用于索引的布尔数组:

from itertools import chain

df = pd.DataFrame(np.random.randint(0, 100, size=(100000, 1)), columns=['state'])
df = df.assign(state=df.state.apply(lambda x: [x]), axis=1)

def wen(df):
    df.state=df.state.astype(str)
    return df.query("state == '[42]'")

%timeit df[np.array(list(chain.from_iterable(df['state'].values))) == 42]  # 14.2 ms
%timeit df[df.state.apply(tuple) == (42,)]                                 # 41.9 ms
%timeit df.loc[df.state.apply(lambda x: x==[42])]                          # 33.9 ms
%timeit wen(df)                                                            # 19.9 ms

更好的是,不要在数据框中使用列表。只需使用常规的int系列。这样可以提高内存效率和性能。

答案 1 :(得分:2)

您可以添加astype(str)

df.state=df.state.astype(str)
df.query("state == '[42]'")
Out[290]: 
  state  action  reward
3  [42]      16      12

答案 2 :(得分:0)

def mVect = udf((arr: Array[Array[Int]]) => arr.transpose.map(_.sum))

另一个解决方案(来自下面的@ user3483203注释):

print df[df.state.apply(tuple) == (42,)]
  state  action  reward
3  [42]  16      12    

但是原始速度要快14%:

df.loc[df.state.apply(lambda x: x==[42])]