「文字列収集」(ランク S) を解いてみた

はじめに

paiza x Qiita のプログラミング問題 の 文字列収集 をトライしました. 正規表現とリストの内包表記の勉強になりました.

末尾に Google Colaboratory へのリンクを貼っています.

ここでは文字列の一致に re.match() を使っていますが, 代わりに startswith() を使った記事も書きましたので, ご参考ください.

「文字列収集」(ランク S) を解いてみた (startswith() を使う)

コード

ライブラリ

正規表現を使います.

1

import re # 正規表現

関数定義

2 つ関数を準備しました. 1 つ目の check_match(s_list, query) は, 読み込んだ文字列のリストとクエリー (文字列) を受け取り, マッチした s_list 内の文字列のインデックスを, リストとして返します.

リストの全部のリストに作用させるために map() を使い, その内部では無名関数 lambda を使って処理を書いてます. クエリーにマッチするかしないかの判定は，re.match() を使っています.

そして得られたリストのうち, 要素が None でないものだけをリスト化し, 返しています (ここも内包表現).

2-1

def check_match(s_list, query):
    """文字列のリスト (s_list) とクエリー (query) を受け取り, マッチしたリストのインデックスを返す
    """
    match_list = list(map(lambda x: re.match(query, x), s_list)) # マッチしたかしないかのリストを返す
    return [i for i, x in enumerate(match_list) if x != None] # マッチしたリストのインデックスを返す

2 つ目の calc_price(p_list, idx_list) は, 価格のリスト p_list とインデックスのリスト idx_list から, 該当するインデックスの価格の合計を計算します. Google Colab のサジェストで一発作成.

2-2

def calc_price(p_list, idx_list):
    """価格のリスト (p_list) とマッチしたインデックスのリスト (idx_list) を受け取り, 合計価格を返す
    """
    return sum([p_list[i] for i in idx_list])

標準入力からデータを読み込む

ここから本体です. 標準入力から指定通りにデータを読んでいきます. 数値は適宜 int() を使って変換しています.

3

line = input().rstrip().split()
n = int(line[0]) # 文字列の数
m = int(line[1]) # クエリーの数

s = [] # 文字列のリスト
p = [] # 価格のリスト
for i in range(n):
    line = input().rstrip().split()
    s.append(line[0])
    p.append(int(line[1]))

q = [] # クエリーのリスト
for i in range(m):
    q.append(input().rstrip())

合計金額を計算して出力する

合計金額を計算して出力します. 上で定義した関数を用いて

1. 文字列がマッチするリストのインデックスを求める
2. 該当するインデックスの価格の合計を求め, 出力する

という処理です.

4

for query in q:
    idx = check_match(s, query)
    print(calc_price(p, idx))

さいごに

これで 100 点を取れたのですが, 一部のテストケースでは秒単位で時間がかかっています. 改善の余地もありそうですが, ひとまずはここまでにします.

リンク

• 文字列収集.ipynb