求助：实现按Z值区分颜色的3D散点图，现有循环代码有误

我想写一个简单的逻辑生成3D散点图，要求根据数据中Z值的大小使用两种不同颜色。我知道自己写的循环有问题，希望有人指出错误所在。

我的代码：

import matplotlib.pyplot as plt
import numpy as np

data = np.loadtxt('Koordinater/Co_sektions.txt')
z = data[:,2]
fig = plt.figure(figsize=(10,7))
ax = plt.axes(projection = "3d")
ax.set_zlim(0, 27000)

d1 = range(0,len(data),1)

for i in d1:
    if z[i] <= 1000:
        ax.scatter3D(data[:,1], data[:,0], data[:,2],color='red')
    else:
        ax.scatter3D(data[:,1], data[:,0], data[:,2])
plt.show()

数据样例：

0  -14450       0
  1583  -14450     902
  1583   14450     902
     0   14450       0
  3166   14450    1704
  3166  -14450    1704
  4749  -14450    2414
  4749   14450    2414
  6332  -14450    3036
  6332   14450    3036
  7915   14450    3576

错误分析 #

你的循环逻辑完全错误：

• 循环遍历单个数据点时，每次调用scatter3D都绘制整个数据集
• 这会导致所有点被重复绘制len(data)次，最后一次绘制的颜色会覆盖之前的所有结果，最终所有点颜色统一，完全达不到分色的目的

修正方案 #

不需要用循环，直接通过布尔索引筛选数据分组绘制，效率更高且逻辑清晰：

import matplotlib.pyplot as plt
import numpy as np

data = np.loadtxt('Koordinater/Co_sektions.txt')
x = data[:, 0]
y = data[:, 1]
z = data[:, 2]

fig = plt.figure(figsize=(10,7))
ax = plt.axes(projection="3d")
ax.set_zlim(0, 27000)

# 筛选Z值<=1000的点，设置为红色
mask_red = z <= 1000
ax.scatter3D(x[mask_red], y[mask_red], z[mask_red], color='red', label='Z ≤ 1000')

# 剩下的点使用默认颜色
mask_other = ~mask_red
ax.scatter3D(x[mask_other], y[mask_other], z[mask_other], label='Z > 1000')

ax.legend()
plt.show()

如果非要用循环（不推荐，效率极低），需要每次只绘制当前单个点：

import matplotlib.pyplot as plt
import numpy as np

data = np.loadtxt('Koordinater/Co_sektions.txt')
fig = plt.figure(figsize=(10,7))
ax = plt.axes(projection="3d")
ax.set_zlim(0, 27000)

for i in range(len(data)):
    x = data[i, 0]
    y = data[i, 1]
    z_val = data[i, 2]
    if z_val <= 1000:
        ax.scatter3D(x, y, z_val, color='red')
    else:
        ax.scatter3D(x, y, z_val)
plt.show()

内容的提问来源于stack exchange，提问作者Mandsberg1