[Ruby|Python]paiza リアルイベント問題セット 文字列収集 (paizaランク S 相当)

今回はpaiza リアルイベント問題セット 「文字列収集」 (paizaランク S 相当)を解説します。

nセットの文字列s_1~s_nと価格p_1~p_n、m個の文字列q_1~q_mが与えられ、各文字列q_iに対して、頭文字がq_iの文字列s_jの価格p_jの合計金額を答える問題です。

文字列収集 (paizaランク S 相当)

※ paiza リアルイベント問題セット より

問題

攻略ポイント

問題を解く流れ

全体の処理を整理する

• 入力データを受け取る
  • 売っている単語の数: 整数n
  • 欲しい頭文字の数: 整数m
  • 売っている単語: word_prices[単語文字列, 単語の価格]
  • 欲しい頭文字: requests
• 売っている単語: word_prices から、ハッシュの価格表price_listを作成する
  • • 価格表price_listの既定値を0で初期化する
      (キーが登録されてないときに0を返す)
    • word_prices 先頭から順に単語wordと価格priceを参照する
      • i を 0 から wordの文字数-1 まで増やしながら繰り返す
        • word の先頭からインデックスi までの文字列を抜き出して key とする
        • price_list のキーに key が未登録なら price_list[key] = price で新規登録
        • price_list のキーに key が登録済なら price_list[key] += price で価格を加算
• 購入希望単語の金額(購入不可なら 0)をresultsに格納する
• 購入希望の購入金額を返す
  （改行で連結した文字列に変換して返す）

処理の流れはこんな感じです。

INPUT1

INPUT2

Rubyで実装してみる

入出力例をコピペしてヒアドキュメントで変数に代入し、データを確認します。正しく受け取れていれば確認用コードは削除します。

INPUT1 = <<~"EOS"
  6 5
  bcac 3
  abcd 14
  abccjg 92
  bcaddgie 2
  abcd 6
  cb 200
  b
  a
  abcd
  gagioheo
  cb
EOS
OUTPUT1 = <<~"EOS"
  5
  112
  20
  0
  200
EOS

INPUT2 = <<~"EOS"
  5 3
  paiza 16
  pizza 1
  paizaio 4
  paizapoh 2
  pizzaio 8
  paiza
  pizza
  p
EOS
OUTPUT2 = <<~"EOS"
  22
  9
  31
EOS

# 確認用コード
p INPUT1
# > "6 5\nbcac 3\nabcd 14\nabccjg 92\nbcaddgie 2\nabcd 6\ncb 200\nb\na\nabcd\ngagioheo\ncb\n"
p OUTPUT1
# > "5\n112\n20\n0\n200\n"
p INPUT2
# > "5 3\npaiza 16\npizza 1\npaizaio 4\npaizapoh 2\npizzaio 8\npaiza\npizza\np\n"
p OUTPUT2
# > "22\n9\n31\n"

続いて問題を解くメソッド（ solve メソッドとします）を変数の下に定義します。

まず、入力データを受け取る処理を書きます。

def solve(input_str)
  # -----------------------------------
  # 入力受け取り
  # -----------------------------------
  # 入力文字列を改行で分割する
  input_lines = input_str.chomp.split("\n")
  # n: 単語と価格の数, m: 購入したい頭文字列の数
  n, m = input_lines.shift.split.map(&:to_i)
  # word_prices: [単語, 価格] の配列
  word_prices = input_lines.shift(n).map do |line|
    word, price = line.split
    [word, price.to_i]
  end
  # requests: 購入したい頭文字の配列
  requests = input_lines.shift(m)

  # 確認用コード
  [n, m, word_prices, requests]
end

# 確認用コード
pp solve(INPUT1)
# > [6, 5,
# >  [["bcac", 3],
# >   ["abcd", 14],
# >   ["abccjg", 92],
# >   ["bcaddgie", 2],
# >   ["abcd", 6],
# >   ["cb", 200]],
# >  ["b", "a", "abcd", "gagioheo", "cb"]]
pp solve(INPUT2)
# > [5, 3,
# >  [["paiza", 16],
# >   ["pizza", 1],
# >   ["paizaio", 4],
# >   ["paizapoh", 2],
# >   ["pizzaio", 8]],
# >  ["paiza", "pizza", "p"]]

入力データを正しく受けとれていれば、単語の価格word_priceを基に、頭文字に対応した価格表price_listを作成します。

※ 計算量は O(N = 10,000 * S_i.length = 100) = 10^6

def solve(input_str)
  # -----------------------------------
  # 入力受け取り
  # -----------------------------------
  # 入力文字列を改行で分割する
  input_lines = input_str.chomp.split("\n")
  # n: 単語と価格の数, m: 購入したい頭文字列の数
  n, m = input_lines.shift.split.map(&:to_i)
  # word_prices: [単語, 価格] の配列
  word_prices = input_lines.shift(n).map do |line|
    word, price = line.split
    [word, price.to_i]
  end
  # requests: 購入したい頭文字の配列
  requests = input_lines.shift(m)

  # -----------------------------------
  # price_list: 価格表を作成する
  # -----------------------------------
  price_list = Hash.new(0)
  word_prices.each do |word, price|
    # word の1文字目から最後まで price を登録する
    0.upto(word.length - 1) do |i|
      # 先頭から index: i までの文字を抜き出す
      key = word[..i]
      if price_list[key].nil?
        # 頭文字 key の価格が未登録なら price で新規登録
        price_list[key] = price
      else
        # 頭文字 key の価格が登録済なら price を加算
        price_list[key] += price
      end
    end
  end

  # 確認用コード
  price_list
end

# 確認用コード
pp solve(INPUT1)
# > {"b"=>5,
# >  "bc"=>5,
# >  "bca"=>5,
# >  "bcac"=>3,
# >  "a"=>112,
# >  "ab"=>112,
# >  "abc"=>112,
# >  "abcd"=>20,
# >  "abcc"=>92,
# >  "abccj"=>92,
# >  "abccjg"=>92,
# >  "bcad"=>2,
# >  "bcadd"=>2,
# >  "bcaddg"=>2,
# >  "bcaddgi"=>2,
# >  "bcaddgie"=>2,
# >  "c"=>200,
# >  "cb"=>200}
pp solve(INPUT2)
# > {"p"=>31,
# >  "pa"=>22,
# >  "pai"=>22,
# >  "paiz"=>22,
# >  "paiza"=>22,
# >  "pi"=>9,
# >  "piz"=>9,
# >  "pizz"=>9,
# >  "pizza"=>9,
# >  "paizai"=>4,
# >  "paizaio"=>4,
# >  "paizap"=>2,
# >  "paizapo"=>2,
# >  "paizapoh"=>2,
# >  "pizzai"=>8,
# >  "pizzaio"=>8}

INPUT1

INPUT2

price_listが完成しましたので、requestsから欲しい頭文字を順に参照して、価格を取得します。

※ 計算量はO(M = 10,000 * 1) = 10^4

def solve(input_str)
  # -----------------------------------
  # 入力受け取り
  # -----------------------------------
  # 入力文字列を改行で分割する
  input_lines = input_str.chomp.split("\n")
  # n: 単語と価格の数, m: 購入したい頭文字列の数
  n, m = input_lines.shift.split.map(&:to_i)
  # word_prices: [単語, 価格] の配列
  word_prices = input_lines.shift(n).map do |line|
    word, price = line.split
    [word, price.to_i]
  end
  # requests: 購入したい頭文字の配列
  requests = input_lines.shift(m)

  # -----------------------------------
  # price_list: 価格表を作成する
  # -----------------------------------
  price_list = Hash.new(0)
  word_prices.each do |word, price|
    # word の1文字目から最後まで price を登録する
    0.upto(word.length - 1) do |i|
      # 先頭から index: i までの文字を抜き出す
      key = word[..i]
      if price_list[key].nil?
        # 頭文字 key の価格が未登録なら price で新規登録
        price_list[key] = price
      else
        # 頭文字 key の価格が登録済なら price を加算
        price_list[key] += price
      end
    end
  end

  # -----------------------------------
  # 希望単語の購入金額(購入不可なら 0)をresultsに格納する
  # -----------------------------------
  results = requests.map { |k| price_list[k] }

  # -----------------------------------
  # 購入希望の頭文字の購入金額を返す
  # （改行で連結した文字列に変換して返す）
  # -----------------------------------
  results.join("\n") << "\n"
end

# 確認用コード
p solve(INPUT1)
# > "5\n112\n20\n0\n200\n"
p solve(INPUT1) == OUTPUT1
# > true
p solve(INPUT2)
# > "22\n9\n31\n"
p solve(INPUT2) == OUTPUT2
# > true

確認用コードを標準入力からのデータ受け取りに変更、標準出力をpメソッドからputsメソッドに変更して提出します。（STDIN.readの入力終了は Ctrl+D)

Ruby 解答例1

def solve(input_str)
  # -----------------------------------
  # 入力受け取り
  # -----------------------------------
  # 入力文字列を改行で分割する
  input_lines = input_str.chomp.split("\n")
  # n: 単語と価格の数, m: 購入したい頭文字列の数
  n, m = input_lines.shift.split.map(&:to_i)
  # word_prices: [単語, 価格] の配列
  word_prices = input_lines.shift(n).map do |line|
    word, price = line.split
    [word, price.to_i]
  end
  # requests: 購入したい頭文字の配列
  requests = input_lines.shift(m)

  # -----------------------------------
  # price_list: 価格表を作成する
  # -----------------------------------
  price_list = Hash.new(0)
  word_prices.each do |word, price|
    # word の1文字目から最後まで price を登録する
    0.upto(word.length - 1) do |i|
      # 先頭から index: i までの文字を抜き出す
      key = word[..i]
      if price_list[key].nil?
        # 頭文字 key の価格が未登録なら price で新規登録
        price_list[key] = price
      else
        # 頭文字 key の価格が登録済なら price を加算
        price_list[key] += price
      end
    end
  end

  # -----------------------------------
  # 希望単語の購入金額(購入不可なら 0)をresultsに格納する
  # -----------------------------------
  results = requests.map { |k| price_list[k] }

  # -----------------------------------
  # 購入希望の頭文字の購入金額を返す
  # （改行で連結した文字列に変換して返す）
  # -----------------------------------
  results.join("\n") << "\n"
end

puts solve(STDIN.read)

Ruby 解答例2（非推奨）

欲しい頭文字request と単語の価格word_price の二重ループを使う方法です。タイムオーバーするかと思いましたが正解してしまいました。

※ 計算量 O(N = 10,000 * M = 10,000) = 10^8

def solve(input_str)
  # -----------------------------------
  # 入力受け取り
  # -----------------------------------
  # 入力文字列を改行で分割する
  input_lines = input_str.chomp.split("\n")
  # n: 単語と価格の数, m: 購入したい頭文字列の数
  n, m = input_lines.shift.split.map(&:to_i)
  # word_prices: [単語, 価格] の配列
  word_prices = input_lines.shift(n).map do |line|
    word, price = line.split
    [word, price.to_i]
  end
  # requests: 購入したい頭文字の配列
  requests = input_lines.shift(m)

  # -----------------------------------
  # 購入希望の頭文字の購入金額を計算する
  # -----------------------------------
  # results: 購入金額の計算結果
  results = []
  # 希望の頭文字を順に参照する
  requests.each do |request|
    # 希望の頭文字の購入金額
    total_price = 0
    # 価格が登録された単語を順に参照する
    word_prices.each do |word, price|
      # 単語の頭文字が希望の頭文字で始まるなら価格を購入金額に加算する
      total_price += price if word.start_with?(request)
    end
    # 希望の頭文字の購入金額を格納する
    results << total_price
  end

  # -----------------------------------
  # 購入希望の頭文字の購入金額を返す
  # （改行で連結した文字列に変換して返す）
  # -----------------------------------
  results.join("\n") << "\n"
end

puts solve(STDIN.read)

Python3 解答例1

Ruby解答例1をPythonに翻訳しました。

def solve(input_str):
    # -----------------------------------
    # 入力受け取り
    # -----------------------------------
    # 入力文字列を改行で分割する
    input_lines = input_str.strip().split("\n")
    # n: 単語と価格の数, m: 購入したい頭文字列の数
    n, m = map(int, input_lines[0].split())
    # word_prices: [単語, 価格] の配列
    word_prices = []
    for line in input_lines[1:-m]:
        word, price = line.split()
        word_prices.append([word, int(price)])
    # requests: 購入したい頭文字の配列
    requests = input_lines[-m:]

    # -----------------------------------
    # price_list: 価格表を作成する
    # -----------------------------------
    price_list = {}
    for word, price in word_prices:
        # word の1文字目から最後まで price を登録する
        for i in range(len(word)):
            key = word[:i+1]
            price_list[key] = price_list.get(key, 0) + price

    # -----------------------------------
    # 希望単語の購入金額(購入不可なら 0)をresultsに格納する
    # -----------------------------------
    results = [price_list.get(key, 0) for key in requests]

    # -----------------------------------
    # 希望単語の購入金額を返す
    # （改行で連結した文字列に変換して返す）
    # -----------------------------------
    return "\n".join(map(str, results)) + "\n"


print(solve(open(0).read()))

Python3 解答例2（非推奨）

Ruby解答例2をPythonに翻訳しました。

def solve(input_str):
    # -----------------------------------
    # 入力受け取り
    # -----------------------------------
    # 入力文字列を改行で分割する
    input_lines = input_str.strip().split("\n")
    # n: 単語と価格の数, m: 購入したい頭文字列の数
    n, m = map(int, input_lines[0].split())
    # word_prices: [単語, 価格] の配列
    word_prices = []
    for line in input_lines[1:-m]:
        word, price = line.split()
        word_prices.append([word, int(price)])
    # requests: 購入したい頭文字の配列
    requests = input_lines[-m:]

    # -----------------------------------
    # 購入希望の頭文字の購入金額を計算する
    # -----------------------------------
    # results: 購入金額の計算結果
    results = []
    # 希望の頭文字を順に参照する
    for request in requests:
        # 希望の頭文字の購入金額
        total_price = 0
        # 価格が登録された単語を順に参照する
        for word, price in word_prices:
            # 単語の頭文字が希望の頭文字で始まるなら価格を購入金額に加算する
            if word.startswith(request):
                total_price += price
        # 希望の頭文字の購入金額を格納する
        results.append(total_price)

    # -----------------------------------
    # 希望単語の購入金額を返す
    # （改行で連結した文字列に変換して返す）
    # -----------------------------------
    return "\n".join(map(str, results)) + "\n"


print(solve(open(0).read()))

解答例の処理速度比較

今回のまとめ

今回は仮想配列（ハッシュ、辞書）を使って計算量を減らす工夫をしてみました。

提出前に、ざっくり計算量を見積もってみると良いかもしれませんね。