Question

我有一个大小为my_array的numpy数组100x20。我想创建一个函数，作为输入接收一个2d numpy数组my_arr和一个索引x，并返回两个数组，其中一个大小为1x20 test_arr，另一个数组为{ {1}} 99x20。向量train_arr将对应于矩阵test_arr的行，其索引为my_arr，而x将包含其余行。我尝试使用屏蔽来遵循解决方案：

train_arr

显然这不符合我的要求。我怎样才能返回一个numpy数组，以及火车和测试阵列？

Answer 1

您可以使用简单索引和numpy.delete：

def split_train_test(my_arr, x):
    return np.delete(my_arr, x, 0), my_arr[x:x+1]

my_arr = np.arange(10).reshape(5,2)

train, test = split_train_test(my_arr, 2)

train
#array([[0, 1],
#       [2, 3],
#       [6, 7],
#       [8, 9]])

test
#array([[4, 5]])

Answer 2

您还可以使用布尔索引作为掩码：

def split_train_test(my_arr, x):

    # define mask
    mask=np.zeros(my_arr.shape[0], dtype=bool)
    mask[x] = True # True only at index x, False elsewhere

    return my_arr[mask, :], my_arr[~mask, :]

示例运行：

test_arr, train_arr = split_train_test(np.random.rand(100, 20), x=10)

print(test_arr.shape, train_arr.shape)
((1L, 20L), (99L, 20L))

修改

如果有人正在寻找需要将多个元素分配给测试数组的一般情况（例如80％-20％分割），x也可以接受数组：

my_arr = np.random.rand(100, 20) x = np.random.choice(np.arange(my_arr.shape[0]), int(my_arr .shape[0]*0.8), replace=False) test_arr, train_arr = split_train_test(my_arr, x) print(test_arr.shape, train_arr.shape) ((80L, 20L), (20L, 20L))

在python中使用掩码拆分numpy数组

2 个答案: