Блог

У меня есть массив, заполненный массивами с использованием numpy split, как показано ниже

 [array([3, 0, 0]),
array([1, 1, 1]),
array([2 , 2, 4]),
array([0, 1, 2]),
array([0, 1, 2]),
array([0, 2, 2])]
 

Я хочу проверить, когда каждый столбец достигнет числа 3 или выше, а затем сохранить этот номер столбца в качестве другого массива или списка. Я также хочу, чтобы он записал номер последнего столбца, если он никогда не превышал 3. Поэтому ответ на этот пример должен быть

 [1, 6, 3]
 

Это означает, что первый столбец достиг числа 3 в строке 1, столбец 2 никогда не достигал числа 3, а столбец 3 достиг числа 3 в строке 3.

Я пробовал несколько вещей, но ни одна из них не сработала.

Вы можете использовать argmax from numpy для достижения этой цели:

 import numpy as np

split_data = [
    np.array([3, 0, 0]),
    np.array([1, 1, 1]),
    np.array([2, 2, 4]),
    np.array([0, 1, 2]),
    np.array([0, 1, 2]),
    np.array([0, 2, 2]),
]
print(f"split_data[0].shape {split_data[0].shape}")
print(f"split_datan{split_data}")

# convert the list of arrays to a single matrix
data = np.stack(split_data)
print(f"data.shape {data.shape}")
print(f"datan{data}")

# create a boolean array where the data is over 3
over3 = data >= 3
print(f"over3n{over3}")

# this finds out the row where 3 was reached
index = np.argmax(over3, axis=0)
print(f"index {index}")

# you start counting from one, so increase that
index = index   1
print(f"index increased {index}")

# this finds out if a column never reached 3
never = np.sum(over3, axis=0)
print(f"never {never}")

# if the sum is 0, then it never went over
# use that to update the index found
index[never == 0] = data.shape[0]
print(f"index {index}")
 

Что приведет к

 split_data[0].shape (3,)
split_data
[array([3, 0, 0]), array([1, 1, 1]), array([2, 2, 4]), array([0, 1, 2]), array([0, 1, 2]), array([0, 2, 2])]
data.shape (6, 3)
data
[[3 0 0]
 [1 1 1]
 [2 2 4]
 [0 1 2]
 [0 1 2]
 [0 2 2]]
over3
[[ True False False]
 [False False False]
 [False False  True]
 [False False False]
 [False False False]
 [False False False]]
index [0 0 2]
index increased [1 1 3]
never [1 0 1]
index [1 6 3]
 

Ваше здоровье!