u++の備忘録

言語処理100本ノック 2020「16. ファイルをN分割する」

問題文

nlp100.github.io

問題の概要

コマンドライン引数でNを受け取る②pandasでデータを読み込み行数を取得する③分割後の各ファイルの行数を計算する④各ファイルを保存するーーの4段階で処理します。

①は「14. 先頭からN行を出力」で解説済み、②は「10. 行数のカウント」そのものです。

③について、X行のファイルをN分割する場合の各ファイルの行数を考えます。XがNで割り切れない場合もあるので(X÷N)を計算した後に、切り上げ・切り下げ・四捨五入などの処理が必要です。今回は切り上げで処理しています。

④では、「loc[]」を用いて行数指定でデータを分割して保存しています。

import sys
import pandas as pd


if len(sys.argv) == 1:
    print('Set arg n, like "python ch02/ans15.py 5"')
else:
    n = int(sys.argv[1])
    df = pd.read_csv('ch02/popular-names.txt', sep='\t', header=None)
    nrow = -(-len(df) // n)

    for i in range(n):
        df.loc[nrow * i:nrow * (i + 1)].to_csv(f'ch02/ans16_{i}', sep='\t', index=False, header=None)