使用条件拆分多维numpy数组

Question

使用条件拆分多维numpy数组

5

我有一个多维的numpy数组。第一个数组表示数据的质量。0代表好的，1代表不太好的。在进行首次检查时，我只想使用好的数据。如何将数组分成两个新的数组？我的想法不起作用：

good_data = [x for x in data[0,:] if x = 1.0]
bad_data = [x for x in data[0,:] if x = 0.0]

这里有一个小例子，说明我的问题:

import numpy as np

flag = np.array([0., 0., 0., 1., 1., 1.])
temp = np.array([300., 310., 320., 300., 222., 333.])
pressure = np.array([1013., 1013., 1013., 900., 900., 900.])
data = np.array([flag, temp, pressure])

good_data = data[0,:][data[0,:] == 1.0]
bad_data  = data[0,:][data[0,:] == 0.0]

print good_data

print语句给出了[1., 1., 1.]。

但我需要的是[[1., 1., 1.], [300., 222., 333.], [900., 900., 900.]]。

- paulchen

好的，看一下我回答的更新。 - milancurcic

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- milancurcic · Accepted Answer

这是您在寻找的内容吗？

good_data = data[0,:][data[0,:] == 1.0]
bad_data  = data[0,:][data[0,:] == 0.0]

这将返回一个numpy.array数组。

或者，你可以按照你的建议操作，但是将结果列表转换为numpy.array：

good_data = np.array([x for x in data[0,:] if x == 1.0])

注意在这里使用比较运算符==而不是赋值运算符=。对于您的特定示例，可以在迭代第一个索引时使用flag == 1来子集化数据。

good_data = [data[n,:][flag == 1] for n in range(data.shape[0])]

如果您真的希望good_data的元素是列表，请在推导式内进行转换：

good_data = [data[n,:][flag == 1].tolist() for n in range(data.shape[0])]

感谢Jaime指出了这个简单的方法：

good_data = data[:, data[0] == 1]