高级切片：给定一个索引列表，在numpy数组中挑选不同的元素。

Question

我正在实现一个决策算法。在daily_choices数组中，每天都有两种水果可以选择，例如：

daily_choices = np.array([['apple','orange'],['strawberry','orange'],['watermelon','apple']])

现在我有一个包含每天我要选择的水果的列表：

decision = [0,1,0] 

我知道一些基本的切片操作，例如daily_choices[:,0]，意思是将第一列切出来，而daily_choices[:,1]则是将第二列切出来。

我想知道是否有一种方法可以通过类似以下的方式切片取出第一行的第一列、第二行的第二列、第三行的第一列：

期望结果

Input  =>  daily_choices[:,[0,1,0]]
Output =>  ['apple', 'orange', 'watermelon']

然而，它没有给我想要的结果

我知道我可以通过使用zip和for循环来实现我想要的结果

daily_decision
daily_decision = []
for choices, index in zip(daily_choices, decision):
    daily_decision.append(choices[index])
daily_decision

但我想知道是否有可能用一行代码来实现。

- Raven Cheuk

---

使用以下代码：daily_decision = [choices[index] for choices, index in zip(daily_choices, decision)] - Mohit Motwani

Answer 1

使用列表推导式

choices = [['apple', 'orange'], ['strawberry', 'orange'], ['watermelon', 'apple']]
decisions = [0, 1, 0] 

daily_decisions = [day[decision] for day, decision in zip(choices, decision)]
print(daily_decisions)

['苹果', '橙子', '西瓜']

使用numpy

这也可以用NumPy的整数数组索引来解决:

import numpy as np
daily_choices = np.array([['apple','orange'],['strawberry','orange'],['watermelon','apple']])
decisions = [0, 1, 0]

daily_decision = daily_choices[range(len(daily_choices)), decisions]
print(daily_decision)

['苹果', '橙子', '西瓜']

- kalehmann

Answer 2

使用纯粹的numpy：

import numpy as np

daily_choices = np.array([['apple', 'orange'],['strawberry', 'orange'],['watermelon', 'apple']])
decision = np.array([0, 1, 0])

n_fruits = 2

fruit_range = np.reshape(np.arange(n_fruits), (-1, n_fruits))
indices = np.reshape(decision, (len(decision), 1)) == fruit_range

daily_choices[indices]

输出：

array(['apple', 'orange', 'watermelon'], dtype='<U10')

- gmds