ARC136 D問題をRuby(NArray)で計算

AtCoder Regular Contest 136 : D - Without Carry

私のレベルだと普段はこの辺の問題は解けないのですが、これに関してはうまくいった（比較的簡単な方法で解けた）のでメモしておきます。なお、公式解説のほうは難しくてまだ理解できていません。

コード

Ruby で多次元配列を簡単・高速に扱うために、 Numo::NArray を使っています。

require 'numo/narray'

L = 6
S = 10**L - 1  # == 999999

n = gets.to_i
a = gets.split.map!(&:to_i)

na = Numo::UInt32.zeros(*Array.new(L, 10)).inplace!
a.each { |ai| na[ai] += 1 }
na.ndim.times { |k| na = na.cumsum(k) }

cnt = a.sum { |ai| na[S - ai] }
cnt -= na[S / 9 * 4]
cnt /= 2

puts cnt

各次元の大きさが 10 の L 次元配列 na を作り、その累積和を計算しています。計算量は O(N + L · 10^L) です。

考え方

簡単のため、最大2桁（L=2）で考えてみます。

以下の入力を例にとります。

input

9
31 41 59 26 53 58 97 93 23

例えば 41 に足して繰り上がりの起きない数は [31, 41, 26, 53, 58, 23] の6つです。というのも、これらは 99 - 41 == 58 よりも 各桁で数字が小さい（または同じ） ためです。

また、例えば「 53 が繰り上がりを起こさないなら [23, 31, 41] も起こさない」ということが言えます。そのためより大きな数値にはその下の情報も合算しておくのが良さそうです。今回の場合は累積和で簡単に作れます。

ある数 s よりも各桁で数字が小さい（または同じ）要素数を表にしてみます。各桁は別々に考えられるので、桁数ぶんの多次元配列を作ればいいです。

#--- aの要素数を各桁の数字別にカウント ---#
Numo::UInt32#shape=[10,10]
[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 1, 0, 0, 1, 0, 0, 0], 
 [0, 1, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 1, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 1, 0, 0, 0, 0, 1, 1], 
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 1, 0, 0, 0, 1, 0, 0]]

#--- その累積和 ---#
Numo::UInt32#shape=[10,10]
[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
 [0, 0, 0, 1, 1, 1, 2, 2, 2, 2], 
 [0, 1, 1, 2, 2, 2, 3, 3, 3, 3], 
 [0, 2, 2, 3, 3, 3, 4, 4, 4, 4], 
 [0, 2, 2, 4, 4, 4, 5, 5, 6, 7], 
 [0, 2, 2, 4, 4, 4, 5, 5, 6, 7], 
 [0, 2, 2, 4, 4, 4, 5, 5, 6, 7], 
 [0, 2, 2, 4, 4, 4, 5, 5, 6, 7], 
 [0, 2, 2, 5, 5, 5, 6, 7, 8, 9]]

この表を使うと、例えば 41 に足して繰り上がりの起きない数の個数は na[9-4, 9-1] == 6 というふうに O(L) で求まります。全要素に対してカウントすると O(NL) です。間に合うか微妙ですが一旦これで行きます。

---

最後に設問に合わせてカウントを補正します。聞かれている個数は i < j の組のみですが、上の数え方だと i == j や i > j も含んでいます。

i == j というのは「自分同士を足して繰り上がりが起きない」ということです。これは全桁が 0〜4 だけの数値が該当します。まずはこの個数を引きます。（全要素に対して判定してもいいですが、既に累積和があるので na[4, 4] を参照すれば一発です。）

i == j を引いた後は、 i < j と i > j というふうに必ず二重カウントしています。なので2で割れば設問通りの個数が求まります。

表へのアクセス高速化

上の考え方で提出したら、いくつかTLEになってしまいました。累積和の計算部分は N に依らず O(L · 10^L) なので、与えられた数を L 桁に分解して多次元配列にアクセスするのが遅いようです。（Rubyなので尚更）

考えた表を見返してみると、各次元の大きさが 10 なので、 na[5, 8] というのは na[58] と同じことです。つまり配列を1次元とみなしてアクセスすれば桁を分解する必要がありません。 na[9-4, 9-1] も na[99-41] と書き直せます。これなら表へのアクセスは O(L) でなく O(1) で済みます。

C言語などであれば、最初から1次元配列として作っておき、累積和の計算時のみ擬似的に多次元配列として扱えばいいと思います。
※ 公式解説のコードは、結果的にこの累積和を計算しているのでしょうか？