CodeIQ：圧縮できるパターンは何通り？

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「圧縮できるパターンは何通り？」（CodeIQ）を参照してください。

問題の概要

問題を引用します。
一列にアルファベットが並んだ文字列があります。
同じ文字が連続するとき、その文字と連続する文字数に変換することにします。
例えば、「AABBBCEEEE」の場合は「A2B3C1E4」のように変換します。
この場合、元の文字列は10文字でしたが、8文字に変換できました。

標準入力から二つの数字 m と n がコンマ区切りで入力されたとき、m 種類のアルファベットを使った n 文字の文字列のうち、元の文字列よりも短くなるパターンが何通りあるかを求めてください。
※m, n は mn ≦65536を満たす値とします。

例えば、2種類、5文字の場合、以下の10通りがあります。
AAAAA → A5
AAAAB → A4B1
AAABB → A3B2
AABBB → A2B3
ABBBB → A1B4
BAAAA → B1A4
BBAAA → B2A3
BBBAA → B3A2
BBBBA → B4A1
BBBBB → B5

私のプログラム

Rubyで解答しています。

#!/usr/bin/ruby

def solve(m,n)
	ret = 0
	(1..m).to_a.repeated_permutation(n){|arr|
		bef = 0
		counts = Array.new(n, 0)
		i = -1
		for a in arr
			if a != bef then
				i += 1
				bef = a
			end

			counts[i] += 1
		end

		wc = 0
		for j in counts
			if j > 0 then wc += 2 end
		end

		if wc < n then
			ret += 1
		end
	}

	return ret
end


#main
while line = gets
	line.strip!
	if line.empty? then next end

	m,n = line.split(",").map{|a| a.to_i}

	p solve(m,n)
end

解説

この問題はそれほど難しくありません。というか、RubyやPythonのように言語の標準ライブラリに重複順列を列挙する機能がある場合、この問題は簡単に解くことができます。

総数が少ない

問題を読んで気になるのが「m, n は mn ≦65536を満たす値」です。
問題文から元の文字列のパターンはmΠn（m個のアルファベットから重複を許してn回選択する重複順列）である事はすぐにわかります。ということはパターンの総数は65536しかないということなので総当たりでやっても時間的に余裕があります。
そしてRubyには重複順列を列挙する機能があり、やたらと速いのでこれを使えば簡単にプログラムを書くことができます。

solve()

まず、重複順列を列挙します（5行目）。問題文ではアルファベットとなっていますがコンピュータにとっては文字も数値も大した違いはありませんので1〜mの数値を使用します。これをrepeated_permutation(n)で重複順列にします。

befはひとつ前の値がなんだったかを覚えておくために使用します。
countsは圧縮後の文字列の数値部分だけを取り出した配列です。「A2B3C1E4」なら[2,3,1,4]というふうになります。配列のサイズは最大でnになる（[1,1,1,…]の場合が最大）のでpush()するオーバヘッドをなくすため静的に確保してしまいます。
iは現在何回文字が変わったかを示します。「AABBBCEEEE」ならAの間は0、Bになった時点で1、Cになった時点で2、Eになった時点で3というふうになります。

9〜16行目の処理は実際に圧縮と同じことをしています。aは重複順列のあるパターンの最初の要素から順に値が入ってくるのでそれをひとつ前の文字と違うかをチェックします（10行目）。違っていたらiを1増やしてbefを現在の文字にします。同じ文字ならcounts[i]を1増やします。

18〜25行目で圧縮後の文字列の長さをチェックします。
wcは圧縮後の文字数をカウントするための変数です。
countsの要素が0で無いならば圧縮後その部分は2文字になります（19〜21行目）。
そしてwcがnより小さかったパターン数を数えます（23〜25行目）。

最後にパターン数を返して終わりです。

雑感

Rubyに重複順列を列挙する機能があったので力任せにやってしまいました。重複順列を列挙する機能が無い場合はメモ化再帰を使って書くことになると思います。この場合、文字列のパターンを作りながらカウントもすることになるので途中で計算を止めることができるため私のコードよりも計算量を減らすことができるはずです。