我想用 numpy 廣播函式重建以下邏輯，例如np.where：如果第一個元素滿足條件，則從 2d 陣列中檢查每行。如果條件為真，則回傳前三個元素作為一行，否則回傳最后三個元素。
我想規避的for回圈形式的簡短MWE：
import numpy as np
array = np.array([
[1, 2, 3, 4],
[1, 2, 4, 2],
[2, 3, 4, 6]
])
new_array = np.zeros((array.shape[0], array.shape[1]-1))
for i, row in enumerate(array):
if row[0] == 1: new_array[i] = row[:3]
else: new_array[i] = row[-3:]
uj5u.com熱心網友回復：
如果你想使用np.where：
import numpy as np
array = np.array([
[1, 2, 3, 4],
[1, 2, 4, 2],
[2, 3, 4, 6]
])
cond = array[:, 0] == 1
np.where(cond[:, None], array[:,:3], array[:,-3:])
輸出：
array([[1, 2, 3],
[1, 2, 4],
[3, 4, 6]])
編輯
稍微簡潔一點的版本：
np.where(array[:, [0]] == 1, array[:,:3], array[:,-3:])
uj5u.com熱心網友回復：
IIUC 你想要這樣的東西：
condition = array[:,0]==1
new_array[condition,:] = array[condition,:3]
new_array[~condition,:] = array[~condition,-3:]
轉載請註明出處，本文鏈接：https://www.uj5u.com/ruanti/448413.html