python - 从数据框中的列创建排列

Question

嗨，我有一个数据框如下：

并想创建一个包含 2 列的数据框：

作家1 作家2

列出了前一首歌曲作者的所有排列：对于歌曲 03 Bonnie & Clyde，作者：Prince、Tupac Shakur、Jay-Z、Tyrone Wrice 和 Kanye West 参与其中。因此，我的数据框应如下所示：

Writer1 Writer2

Prince  Tupac Shakur

Prince  Jay-Z

Prince  Tyrone Wrice

Prince  Kanye West

Tupac S Jay-Z

Tupac S Tyrone Wrice

Tupac S Kanye West

Jay-Z   Tyrone Wrice

Jay-Z   Kanye West

Tyrone  Kanye West

知道我该怎么做吗？

标签： pythondataframepermutation

Answer 1

这是使用的一种方法itertools.combinations：

import itertools
import pandas as pd

def get_combinations(df, song_name):
    """
    Get a dataframe of all two-writer combinations for a given song.

    :param df: dataframe containing all artists, songs and writers
    :param song_name: name of song 
    :returns: dataframe with cols 'Writer1', 'Writer2' of all two writer combinations for the given song
    """
    song_frame = df[df['Song'] == song_name]
    combinations_df = pd.DataFrame(list(itertools.combinations(song_frame['Writer'].unique(), 2)), 
                                   columns=['Writer1', 'Writer2'])
    return combinations_df

combinations_df = get_combinations(df, '03 Bonnie & Clyde')

请注意，这假设您的数据采用 Pandas 数据框的形式。您可以轻松地从文本文件或 csv 中读取数据，或创建如下所示的文件进行测试：

import numpy as np
df = pd.DataFrame({'Artist':np.repeat('Jay-Z',5).tolist() + ['David Bowie'] * 2 + ['List of the X Factor finalists'] * 2,
                   'Song':np.repeat('03 Bonnie & Clyde',5).tolist() + ['Heroes'] * 4,
                   'Writer':['Prince', 'Tupac Shakur',
                             'Jaz-Z', 'Tyrone Wrice',
                             'Kanye West'] + ['David Bowie', 'Brian Eno'] * 2})

如果您想有效地将​​其应用于整个数据框，请考虑：

def combinations_per_group(group):
    """Return combinations of writers after grouping by song."""     
    group_combs = pd.DataFrame(list(itertools.combinations(group['Writer'].unique(),2)),
                               columns=['Writer1', 'Writer2'])
combinations_df = df.groupby(['Song']).apply(combinations_per_group)\
                    .reset_index()\
                    .drop('level_1', axis=1)

这将返回一个数据框，其中歌曲作为索引，所需的列给出每首歌曲的所有作者组合。