Подразделение ndarray на массив категорий

Вы можете использовать словарь для сопоставления массивов с их соответствующими категориями и сбора результатов в конце. Преимущество использования словаря таким образом заключается в том, что он способен обрабатывать категории, которые также выглядят не по порядку.

input_list = [
    [1,1,1],
    [2,2,2],
    [3,3,3],
    [4,4,4]
] 

category = [1, 2, 3, 2]

from collections import defaultdict

temp_dict = defaultdict(list)
#Iterate through each array of input list, and append to temp_dict
for i, arr in zip(category, input_list):
    temp_dict[i].append(arr)

result = list(temp_dict.values())
print(result)
#[[[1, 1, 1]], [[2, 2, 2], [4, 4, 4]], [[3, 3, 3]]]

При желании, если порядок также важен, используйте сортировку по словарным ключам.

result = [x[1] for x in sorted(temp_dict.items())]
print(result)