时间戳

首页 > 解决方案 > 从 Pandas DataFrame 中提取样本，保持所有相同类型的值

python - 从 Pandas DataFrame 中提取样本，保持所有相同类型的值

问题描述

我有一个巨大的 CSV，它是这样制作的：

type, value
A     1
B     4
C     6
A     25
D     5
B     7

由于要处理的行太多，我想取一个样本，但这个样本的特殊性必须如下：必须抽取所有相同类型的行。

我开始随机抽取行样本：

num_lines = sum(1 for line in open('file.csv') - 1
sample_lines = int(num_lines * 0.01)
skip = sorted(random.sample(range(num_lines), num_lines - sample_lines))
df = pd.read_csv('file.csv', sep=';', skiprows=skip)

但这给了我一个随机的行样本。我想获得的是类型的随机样本。

我有一个大致的过程：

  1. 在 Pandas DataFrame 中导入整个 CSV
  2. 生成要提取的（随机）类型列表（例如 [A, B]）
  3. 仅从 DataFrame 中提取类型为“A”或“B”的行

结果应该是这样的：

type  value
A     1
B     4
A     25
B     7

感谢您提供的任何帮助。

标签： pythonpandas

解决方案

这是一个正确的方法吗？

首先，通过从 CSV 导入 DataFrame 来创建它。然后，创建包含所有可能类型列表的数组，并仅选择其中的 n 个（随机）。最后，保存一个仅包含这 n 个类型的新 DataFrame（但包含与它们相关的所有数据）。

n = 10
df = pd.read_csv('file.csv', sep=';')
random_types = np.random.choice(df.type.unique(), n)
m = df['type'].isin(random_types)
df_sample = df.loc[m]

但是，这种方法的缺点是必须将整个 CSV 加载到内存中。

完整示例

import pandas as pd
import numpy as np

np.random.seed(400)

data = '''\
type value
A     1
B     4
C     6
A     25
D     5
B     7'''

fileobj = pd.compat.StringIO(data)
df = pd.read_csv(fileobj, sep='\s+')
n = 2
random_types = np.random.choice(df.type.unique(), n)
print(df.loc[df['type'].isin(random_types)])

回报：

  type  value
0    A      1
3    A     25
4    D      5

推荐阅读