numpy - How to make a boolean index in numpy for multiple integer types without a loop

Question

I would like to create a boolean index of an array that has true values whenever any integers in a given list appear in the array. Right now I do this by looping through the list of test integers and constructing a seperate boolean mask for each which I bitwise or together like this:

boolean_mask = original_array == list_of_codes[0]
for code in list_of_codes[1::]:
    boolean_mask = boolean_mask | (original_array == code)

Is there any shorthand numpy notation to do the equivalent without the loop?

Answer 1

你有np.in1d：

boolean_mask = np.in1d(original_array, list_of_codes)

应该这样做。请注意，np.in1d这两个数组都变平了，所以如果你original_array是多维的，你将不得不重新塑造它，例如：

boolean_mask = boolean_mask.reshape(original_array.shape)

Answer 2

我认为广播会让你写：

boolean_mask = (original_array == list_of_codes)

如果这是错误的，请给出一个小测试示例，以便我可以确保我知道您要做什么。