main.py 文件源码

def getBalancedSample(y, seed, test_sz=1000):
    if y.shape[0] == test_sz:
        return np.arange(test_sz)
    else:
        stratSplit = StratifiedShuffleSplit(
            y, 1, test_size=test_sz, random_state=seed
        )
        for _, test_idx in stratSplit:
            idx = test_idx
            break
        return idx