python - 使用行分隔符转换数据框

Question

我制作了一个接受数据框作为输入的函数：

a = {"string": ['xxx', 'yyy'], "array": [[1,2,3,4,5,6,1,2,3,6,6,2,2,3,5,6], [2,6,6]]}
df = pd.DataFrame(a)

    string  array
0   xxx [1, 2, 3, 4, 5, 6, 1, 2, 3, 6, 6, 2, 2, 3, 5, 6]
1   yyy [2, 6, 6]

并返回一个数据帧，其中某个分隔符编号（在示例中为 6）是传递的参数：

    string  array
0   xxx [1, 2, 3, 4, 5, 6]
1   xxx [1, 2, 3, 6]
2   xxx [6]
3   xxx [2, 2, 3, 5, 6]
4   yyy [2, 6]
5   yyy [6]

这是我得到的：

def df_conversion(df, sep=None):
    data = {}
    idx = []
    
    for i in range(df.shape[0]):       
        key = df['string'].iloc[i]
        value = df['array'].iloc[i]

        spl = [[]]
        for item in value:
            if item == sep:
                spl[-1].append(item)
                idx.append(key)
                spl.append([])
            else:
                spl[-1].append(item)

        del spl[-1]
        if i == 0: spl_0 = spl
        if i == 1: spl_0.extend(spl)

    data['string'] = idx
    data['array'] = spl_0

    return pd.DataFrame(data)

df_conversion(df, 6)

如何简化功能并使其更加通用？如何使功能更快？谢谢。

标签： pythonpandasdataframe

Answer 1

np.split()您可以使用and简洁地做到这一点df.explode()：

sep = 6
df.array = df.array.apply(lambda a:
    np.split(a, 1 + np.where(np.array(a) == sep)[0][:-1]))

df = df.set_index('string').explode('array').reset_index()

#   string               array
# 0    xxx  [1, 2, 3, 4, 5, 6]
# 1    xxx        [1, 2, 3, 6]
# 2    xxx                 [6]
# 3    xxx     [2, 2, 3, 5, 6]
# 4    yyy              [2, 6]
# 5    yyy                 [6]

---

np.split()和的解释np.where()

我们np.where()用来查找 的索引sep：

a = [1, 2, 3, 4, 5, 6, 1, 2, 3, 6, 6, 2, 2, 3, 5, 6]
sep = 6
np.where(np.array(a) == sep)[0]

# array([ 5,  9, 10, 15])

但是，在每个索引之后np.split()进行拆分，这会放在每个拆分的开头：sep

np.split(a, np.where(np.array(a) == sep)[0])

# [array([1, 2, 3, 4, 5]),
#  array([6, 1, 2, 3]),
#  array([6]),
#  array([6, 2, 2, 3, 5]),
#  array([6])]

相反，OP 希望在每个索引之前拆分以保留sep在每个拆分的末尾，因此我们移动拆分索引 ( 1 +) 并删除将不再存在的最后一个拆分索引 ( [:-1])：

np.split(a, 1 + np.where(np.array(a) == sep)[0][:-1])

# [array([1, 2, 3, 4, 5, 6]),
#  array([1, 2, 3, 6]),
#  array([6]),
#  array([2, 2, 3, 5, 6])]