Python 获取数组中等于给定值大小的值

  • 引言
  • 正文
  • 基础方法
  • 方法1
  • 方法2------使用 np.where() 函数
  • 方法3------数组操作
  • 示例1
  • 示例 2


引言

可能很多小伙伴看到这个标题会觉得很绕口，甚至可能会觉得这属于脱裤子放屁，多此一举，但是这里请大家耐心往下看，这篇文章将介绍一个在实际应用中比较常见的数据处理的示例。

正文

基础方法

比如，这里我们有一个 numpy 数组。

[1 2 3 4 5 6 7 8 9 0]

我们想要获取到它中间等于 5 的值。那么该怎么操作呢？

方法1

遍历整个数组进行查找并打印

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 0])
num = 5
for element in arr:
    if element == num:
        print(element)
"""
result:
5
"""

当然，如果加上 enumerate 我们也可以顺带得到其对应的索引值。如下：

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 0])
num = 5
for index, element in enumerate(arr):
    if element == num:
        print(index, element)
"""
result:
4 5
"""

虽然我们得到了结果，但是显然这并不 Pythonic。

方法2------使用 np.where() 函数

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 0])
num = 5
result = arr[np.where(arr == num)]
print(result)
"""
result:
[5]
"""

相比于之前，我们得到的是一个数组，而不再是一个数罢了。至于获取对应的索引值，只需要打印 np.where(arr == num) 即可，如下：

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 0])
num = 5
index = np.where(arr == num)
print(index)
"""
result:
(array([4], dtype=int64),)
"""

但是，这仍旧不是最简单的方式

方法3------数组操作

import numpy as np

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 0])
num = 5
result = arr[arr == num]
print(result)
"""
result:
[5]
"""

可以看到，我们也得到了想要的结果，看到这里，感觉还是在扯，别着急，看一个下面的例子。

示例1

在现实的例子中，往往我们的数组元素本身并不是一些整数，而是小数，比如：

[1.4465446, 2.5445116, 3.654451, 4.156154964, 5.516515, 6.1264847, 
7.126448, 8.1264984, 9.16154654, 0.54881651]

这时，我们想要获取到它们谁更加近似等于我们给定的数据 5 该怎么办呢？

import numpy as np

arr = np.array([1.4465446, 2.5445116, 3.654451, 4.156154964, 5.416515, 6.1264847, 7.126448, 8.1264984, 9.16154654, 0.54881651])
num = 5
result = arr[np.round(arr, 0) == num]
print(result)
"""
result:
[5.416515]
"""

到这里，是不是体会到了标题真正的含义了？有人也许会说，你这里如果是 5.5 几几几就找不到值了，然而这种担心是多余的，因为实际中，往往会存在很多个接近 5 的小数，总会有值的，如果没有，那说明确实也没有数字接近 5。

示例 2

更经典的一个示例是，如果我们有两组一一对应的数据，比如：

1 1.4465446
2 2.5445116
3 3.654451
4 4.156154964
5 5.416515
6 5.123546
7 5.461682
8 5.1565446
9 5.364588979
10 5.7845151

我们想要获取后面一列数字中接近 5 的值所对应的第一列数字的值，但是我们发现后面一列数字中有很多值满足条件，那么如何找出前一列中唯一对应的数字呢？

import numpy as np

index = np.arange(1, 11)
arr = np.array([1.4465446, 2.5445116, 3.654451, 4.156154964, 5.416515, 5.123546, 5.461682, 5.1565446, 5.364588979, 5.7845151])
num = 5
result = index[int(np.round(np.mean(arr[np.round(arr, 0) == num]), 0))]
print(result)
"""
result:
6
"""

至此，我们完全说清楚了这个问题。