【メモ】pandasで重複行の組み合わせを抽出

業務でかなりカラムの多いエクセルから重複したペアを抽出してペア毎に対処する必要があるシーンに出会ったので、重複行の組み合わせを抽出するスクリプトを作成しました。単純に重複を削除するのではなく、「1行目と4行目が一致していること」がわかるような方法が意外と少ないですね。

実行イメージ

〜〜〜入力となるファイル(csv, tsv, Excel)〜〜〜
a,b,a,b,a
b,a,b,a,b
a,b,a,b,a
a,a,a,a,a
b,b,b,b,b
b,b,b,b,b

〜〜〜出力イメージ〜〜〜
### pair list ###
[1, 3]
[2]
[4]
[5, 6]

実装について

すべてのカラムを使ってgroup_byして、groupsを取り出すと重複しているものは同じグループに、ユニークなものは単一の値のグループになります。

df.groupby(list(df.columns)).groups

それだけで良いのですが、indexの最小値でソートする処理いれているのでソースはちょっとごちゃごちゃしちゃいます。

カラムはかなり多いけどレコード数は少ないエクセルとかでやる予定なので、性能度外視。100件くらいならなんとかなると思いますが、数万以上はやめておいたほうが良いと思います。

ちなみに家でエクセルファイルを扱えないので動作確認はcsv, tsvでしかしていません。

import sys
import os
import pandas as pd
import argparse


def parse_args():
    parser = argparse.ArgumentParser(
        description='重複している行の組み合わせを出力する.')
    parser.add_argument('-i', '--input_path', type=str, default='dup.csv',
                        help='input file path')
    parser.add_argument('-b', '--base_index', type=int, default=1,
                        help='value to add line number.')
    parser.add_argument('-t', '--input_type', type=str, default='csv',
                        choices=['csv', 'tsv', 'xlsx'],
                        help='kind of input_file. csv or tsv or xlsx')
    args = parser.parse_args()

    return args


def make_pair_list(df):
    # 全カラムの組み合わせでgroup抽出
    group_dict = df.groupby(list(df.columns)).groups

    # 元のindexが小さい順に並び替える
    tmp_list = list(group_dict.values())
    tmp_df = pd.DataFrame(tmp_list)
    tmp_cols = ['col{}'.format(i) for i in tmp_df.columns]
    tmp_df.columns = tmp_cols
    sort_ids = tmp_df[tmp_cols[0]].sort_values().index

    group_list = [tmp_list[i] for i in sort_ids]
    # listにキャストして返却したほうが見やすい
    return [list(item) for item in group_list]


def main():
    args = parse_args()
    input_path = args.input_path
    base_index = args.base_index
    input_type = args.input_type

    if not os.path.isfile(input_path):
        print('file not found. input_path is ;{}'.format(input_path))
        sys.exit(1)
    df = __read_input_file(input_path, input_type)
    df.index = df.index + base_index
    if df.duplicated().any():
        pair_list = make_pair_list(df)
        print('### pair list ###')
        for pair in pair_list:
            print(pair)
    else:
        print('there are no pair. all lines are unique')


def __read_input_file(input_path, input_type):
    if input_type == 'csv':
        sep = ','
    else:
        sep = '\t'
    if input_type == 'xlsx':
        df = pd.read_excel(input_path, header=None, sheet_name=0)
    else:
        df = pd.read_csv(input_path, header=None, sep=sep)
    return df


if __name__ == '__main__':
    main()