在train_test_split之后运行proportions_ztest

Question

我得到了一个数据集，在将其拆分成训练数据和测试数据后，我尝试通过train_test_split对y_train和y_test运行proportion_ztest

(test_stat, p_value) = proportions_z_test(y_train, y_test, alternative='two-sided')

但是Python一直抛出ValueError: operands could not be broadast together with shapes (4254,) (1123,)。

我的Y目标变量是二进制的(类0和类1)

有没有办法直接将y_train和y_test添加到proportion_ztest调用中，就像上面的代码一样，或者，在此之前，我必须计算每个数据集中的所有1个类和观察值的总数(y_train和y_test)，并使用np.arrays编写代码，如下所示：

success = [123, 359]
TotalObs = [2500, 2500]
(test_stat, p_value) = proportions_z_test(success, TotalObs, alternative='two-sided')

我正在尝试寻找一种解决方案，而不是向train_test_split调用中添加stratify参数。

任何帮助都会被感谢。

泰！

Answer 1

如果随机变量已经编码为二进制0，1，那么sum将计算成功的次数。

因此，下面的方法应该适用于相等比例的两个样本测试

(test_stat, p_value) = proportions_z_test([y_train.sum(), y_test.sum()], 
                                          [len(y_train), len(y_test)], 
                                          alternative='two-sided')