u++の備忘録

言語処理100本ノック 2020「14. 先頭からN行を出力」

問題文

nlp100.github.io

問題の概要

コマンドライン引数でNを受け取る②pandasでデータを読み込み上位N件を表示するーーの2段階で処理します。

Pythonにて実行時に「python (ファイル名) hoge」として与えたコマンドライン引数は「sys.argv」にリストとして格納されています。例えば「python fuga.py hoge」を実行した場合は「sys.argv」は「[fuga.py, hoge]」です。リストの0番目の要素がファイル名、1番目の要素が一つ目のコマンドライン引数です。

コマンドライン引数でNを指定し忘れた場合には、メッセージを表示するような条件分岐を設けました。

pandasでデータを読み込み上位N件を表示するには「head(N)」が利用できます。

import sys
import pandas as pd


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