Ruby と Python と numpy で解く AtCoder ABC054 B 行列演算

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

AtCoder Beginner Contest B - Template Matching
Difficulty: 828

今回のテーマ、行列演算

2次元配列ですが、行列のライブラリを使用すると演算が簡単になることがあります。

Ruby

ruby.rb

require 'matrix'

class Matrix
  # v2.3 では代入不可
  def []=(i, j, x)
    @rows[i][j]=x
  end
end

n, m = gets.split.map(&:to_i)
rows = Array.new(n + m){gets.chomp.chars.map{|c| (c == '#') ? 1 : 0}}
a = Matrix.rows(rows.shift(n))
b = Matrix.rows(rows)
(n - m + 1).times do |i|
  (n - m + 1).times do |j|
    if a.minor(i, m, j, m) == b
      puts "Yes"
      exit
    end
  end
end
puts "No"

matrix.rb

require 'matrix'

class Matrix
  # v2.3 では代入不可
  def []=(i, j, x)
    @rows[i][j]=x
  end
end

require 'matrix'で行列のライブラリを呼び出します。
ローカル環境のv2.7.1では行列の要素に代入できるのですが、AtCoder環境のv2.3.3では代入できないのでメソッドを追加しています。
標準ライブラリに対してもこういうことができるのがRubyの面白い所でもあります。

minor.rb

    if a.minor(i, m, j, m) == b

minorで部分行列を取得し行列の比較を行っています。

error.rb

(n - m).times do |i|
  (n - m).times do |j|

(n - m + 1).times do |i|
  (n - m + 1).times do |j|

1 x 1の行列の時、行列の比較がうまくいかないので、それに対応した処理が入っています。(n - m)は間違いで(n - m + 1)が正解でした。

追記
2次元配列版については、コメント欄を参照願います。

Python

import numpy

n, m = map(int, input().split())
a = numpy.zeros([n, n], dtype=int)
b = numpy.zeros([m, m], dtype=int)
for i in range(n):
    s = input()
    for j in range(n):
        if s[j] == '#':
            a[i, j] = 1
for i in range(m):
    s = input()
    for j in range(m):
        if s[j] == '#':
            b[i, j] = 1
for i in range(n - m + 1):
    for j in range(n - m + 1):
        if numpy.all(a[i: i + m, j: j + m] == b):
            print("Yes")
            exit()
print("No")

all.py

        if numpy.all(a[i: i + m, j: j + m] == b):

numpyの行列の比較は、例えば[[True, True], [True, True]]を返しますので、all関数を用いて全てTrueかどうかを調べています。

	Ruby (Matrix)	Ruby (Array)	Python (numpy)
コード長 (Byte)	420	288	534
実行時間 (ms)	22	10	166
メモリ (KB)	3196	1788	12512

まとめ

• ABC 054 B を解いた
• Ruby に詳しくなった
• Python に詳しくなった
• numpy に詳しくなった