详解NumPy中np.where() 的两种神奇用法

import numpy as np

# 创建示例数组
temperatures = np.array([22, 28, 15, 32, 18, 25])

# 标记高温和低温
result = np.where(temperatures > 25, "高温", "舒适")
print(result)
# 输出：['舒适' '高温' '舒适' '高温' '舒适' '舒适']

scores = np.array([75, 92, 58, 81, 45, 67, 88])

# 根据分数分类
grade = np.where(scores >= 90, "A",
          np.where(scores >= 80, "B",
            np.where(scores >= 70, "C",
              np.where(scores >= 60, "D", "F"))))

print(grade)
# 输出：['C' 'A' 'F' 'B' 'F' 'D' 'B']

data = np.array([12, 25, 7, 18, 30, 5, 22])

# 组合条件：大于10且小于20
result = np.where((data > 10) & (data < 20), data, 0)
print(result)  # 输出：[12  0  0 18  0  0  0]

# 使用 | 表示 OR 条件
result = np.where((data < 10) | (data > 20), data, -1)
print(result)  # 输出：[ -1  25   7  -1  30   5  22]

arr = np.array([0, 5, 0, 8, 0, 3, 0])

# 找到非零元素的索引
non_zero_indices = np.where(arr != 0)
print(non_zero_indices)  # 输出：(array([1, 3, 5]),)

# 提取非零值
print(arr[non_zero_indices])  # 输出：[5 8 3]

matrix = np.array([[1, 0, 4],
                   [0, 5, 0],
                   [7, 0, 9]])

# 找到值大于3的元素位置
rows, cols = np.where(matrix > 3)

print("行索引:", rows)    # 输出：[0 1 2 2]
print("列索引:", cols)    # 输出：[2 1 0 2]
print("对应值:", matrix[rows, cols])  # 输出：[4 5 7 9]

# 创建一个简单的图像矩阵 (5x5)
image = np.array([[120, 130, 40, 200, 210],
                  [30, 145, 255, 180, 10],
                  [220, 25, 30, 190, 200],
                  [100, 110, 120, 130, 140],
                  [50, 60, 70, 80, 90]])

# 找到高光区域（值>200）
highlight_rows, highlight_cols = np.where(image > 200)

print("高光像素位置:")
for r, c in zip(highlight_rows, highlight_cols):
    print(f"({r}, {c}) - 值: {image[r, c]}")

# 输出：
# (0, 3) - 值: 200
# (0, 4) - 值: 210
# (1, 2) - 值: 255
# (2, 0) - 值: 220
# (2, 3) - 值: 190 -> 注意：190不大于200，实际应为 (2, 4): 200
# 更正：矩阵中 (2,4) 是200，所以应包含

# 二维条件与一维值组合
condition_2d = np.array([[True, False], [False, True]])
result = np.where(condition_2d, [10, 20], 0)

print(result)
# 输出：
# [[10  0]
#  [ 0 20]]

data = np.array([5, 12, 8, 15, 3, 10])

# 将小于10的值替换为0
data[np.where(data < 10)] = 0
print(data)  # 输出：[ 0 12  0 15  0 10]

# 创建3x3x3数组
cube = np.random.randint(0, 10, (3, 3, 3))

# 找到所有大于8的元素
indices = np.where(cube > 8)

# 输出三维索引
print("维度0:", indices[0])
print("维度1:", indices[1])
print("维度2:", indices[2])

# 访问这些元素
print("满足条件的值:", cube[indices])

import time

large_array = np.random.rand(10**6)

# 使用循环
start = time.time()
result_loop = [x*2 if x > 0.5 else x/2 for x in large_array]
print("循环耗时:", time.time() - start)

# 使用 np.where
start = time.time()
result_np = np.where(large_array > 0.5, large_array*2, large_array/2)
print("np.where耗时:", time.time() - start)