如何获得指数的分层随机样本？

Question

我有一个包含两个值（例如 A 和 B）的数组（pd.Series）。

y = pd.Series(['A','B','A','A','B','B','A','B','A','B','B'])


0 A
1 B
2 A
3 A
4 B
5 B
6 A
7 B
8 A
9 B
10 B

我想从系列中获取索引的随机样本，但一半的索引必须与 A 对应，另一半必须与 B 对应。

例如

get_random_stratified_sample_of_indices(y=y, n=4)

[0, 1, 2, 4]

索引 0 和 2 对应于 A 的索引，索引 1 和 4 对应于 B 的索引。

另一个例子

get_random_stratified_sample_of_indices(y=y, n=6)

[1, 4, 5, 0, 2, 3]

返回的索引列表的顺序无关紧要，但我需要它甚至在 y 数组中的 A 和 B 的索引之间进行拆分。

我的计划是首先查看 A 的索引，然后对索引进行随机抽样 (size=n/2)。 然后对 B 重复。

Answer 1

您可以使用groupby.sample ：

N = 4

idx = (y
  .index.to_series()
  .groupby(y)
  .sample(n=N//len(y.unique()))
  .to_list()
 )

Output： [3, 8, 10, 1]

查看：

3     A
8     A
10    B
1     B
dtype: object

Answer 2

这是一种方法：

def get_random_stratified_sample_of_indices(s, n):
    mask = s == 'A'
    s1 = s[mask]
    s2 = s[~mask]
    m1 = n // 2
    m2 = m1 if n % 2 == 0 else m1 + 1
    i1 = s1.sample(m1).index.to_list()
    i2 = s2.sample(m2).index.to_list()
    return i1 + i2

可以这样使用：

y = pd.Series(['A','B','A','A','B','B','A','B','A','B','B'])
i = get_random_stratified_sample_of_indices(y, 5)
print(i)
print()
print(y[i])

结果：

[6, 2, 7, 10, 5]

6     A
2     A
7     B
10    B
5     B

Answer 3

我认为您可以使用train_test_split中的 train_test_split 定义其stratify参数。

from sklearn.model_selection import train_test_split
import pandas as pd

y = (
    pd.Series(["A", "B", "A", "A", "B", "B", "A", "B", "A", "B", "B"])
    .T.to_frame("col")
    .assign(i=lambda xdf: xdf.index)
)
print(y)
# Prints:
#
#    col   i
# 0    A   0
# 1    B   1
# 2    A   2
# 3    A   3
# 4    B   4
# 5    B   5
# 6    A   6
# 7    B   7
# 8    A   8
# 9    B   9
# 10   B  10
print('\n')

# ===== Actual solution =====================================
a, b = train_test_split(y, test_size=0.5, stratify=y["col"])
# ===========================================================
print(a)
# Prints:
#
#    col   i
# 10   B  10
# 6    A   6
# 7    B   7
# 8    A   8
# 4    B   4

print('\n')
print(b)
# Prints:
#
#   col  i
# 3   A  3
# 9   B  9
# 2   A  2
# 1   B  1
# 5   B  5
# 0   A  0