Python 如何将数组分成几个部分

发布于08月07日

我有一个长度约为12000的array.另外，我在DataFrame中有一列提供了类似2,3,7,3,2,7...的值.该列的长度为48，这些值的总和为36.

我想要分布值，这意味着12000长度的数组是按特定的每个值分布的.例如，该列中的第一个值(=2)获得它自己的12000*(2/36)数组(可能是[0.3，0.6，0.3])，第二个值(=3)获得它的12000*(3/36)数组，它的值在第一个值(类似于[0.5，0.1，0.9，0.4])之后继续，依此类推.

推荐答案

import pandas as pd
import numpy as np


# mock some data
a = np.random.random(12000)
df = pd.DataFrame({'col': np.random.randint(1, 5, 48)})

indices = (len(a) * df.col.to_numpy() / sum(df.col)).cumsum()
indices = np.concatenate(([0], indices)).round().astype(int)
res = []
for s, e in zip(indices[:-1], indices[1:]):
    res.append(a[round(s):round(e)])

# some tests
target_pcts = df.col.to_numpy() / sum(df.col)
realized_pcts = np.array([len(sl) / len(a) for sl in res])
diffs = target_pcts / realized_pcts
assert 0.99 < np.min(diffs) and np.max(diffs) < 1.01
assert all(np.concatenate([*res]) == a)