[AtCoder Nx] ABC087B - Coins [Elixir Livebook]

はじめに

AtCoder で Elixir の Nx が使えるようになったので再入門するシリーズです

以前、 Nx なしで解いたのはこちら

今回実装したノートブックはこちら

問題

入出力例

以下のようになれば OK です

"""
2
2
2
100
"""
|> Main.solve()
|> then(&(&1 == 2))

"""
5
1
0
150
"""
|> Main.solve()
|> then(&(&1 == 0))

"""
30
40
50
6000
"""
|> Main.solve()
|> then(&(&1 == 213))

Nx なしの回答

内包表現を用いて条件を満たす組み合わせを探索し、その数を出しています

defmodule Main do
  def main do
    :stdio
    |> IO.read(:all)
    |> solve()
    |> IO.puts()
  end

  defp split_lines(lines) do
    lines
    |> String.trim()
    |> String.split("\n")
  end

  def solve(input) do
    [num_500, num_100, num_50, total] =
      input
      |> split_lines()
      |> Enum.map(&String.to_integer/1)

    total_50 = total / 50

    for p_500 <- 0..num_500,
        p_100 <- 0..num_100,
        p_50 <- 0..num_50,
        10 * p_500 + 2 * p_100 + p_50 == total_50 do
      {p_500, p_100, p_50}
    end
    |> Enum.count()
  end
end

実行時間は 804 msec でした

Nx ありの回答

堂々と全ての組み合わせを計算し、合計金額が X 円になった組み合わせの数を出します

今回の場合、3種類の硬貨の全ての組み合わせを計算するため、例えば枚数が以下のような場合を考えます

• 500 円玉 2 枚
• 100 円玉 2 枚
• 50 円玉 2 枚

ここから各硬貨を何枚かずつ使う組み合わせを考えます

まず 500 円玉と 100 円玉だけで考えてみます
組み合わせは以下のように 9 通り考えられます

• 500 * 0 + 100 * 0
• 500 * 0 + 100 * 1
• 500 * 0 + 100 * 2
• 500 * 1 + 100 * 0
• 500 * 1 + 100 * 1
• 500 * 1 + 100 * 2
• 500 * 2 + 100 * 0
• 500 * 2 + 100 * 1
• 500 * 2 + 100 * 2

これは 0 から 2 までの整数の組み合わせを考えているのと同じです

Nx.add() を使うと、全ての和の組み合わせを計算できます

[0, 1, 2]　 と [0, 1, 2] の全ての和の組み合わせは以下のように計算します

\begin{bmatrix}
0 & 1 & 2
\end{bmatrix}
^\top
+
\begin{bmatrix}
0 & 1 & 2
\end{bmatrix}
=
\begin{bmatrix}
0 \\
1 \\
2
\end{bmatrix}
+
\begin{bmatrix}
0 & 1 & 2
\end{bmatrix}
=
\begin{bmatrix}
0 && 1 && 2 \\
1 && 2 && 3 \\
2 && 3 && 4
\end{bmatrix}

ここでは以下のようなイメージの計算をしています

-	0	1	2
0	0+0	0+1	0+2
1	1+0	1+1	1+2
2	2+0	2+1	2+2

これは次元が増えても同じように計算できます

つまり、片方の次元を増やして転置して足せば良いだけです

500 円玉の枚数を n としたとき、使った枚数毎の金額を表す行列は以下のコードで表せます

Nx.iota(x) = [0, 1, 2, ..., x - 1]

{n + 1}
|> Nx.iota()
|> Nx.multiply(500)

これらを利用して、最終的なコードは以下のようになります

defmodule Main do
  import Nx.Defn

  def main do
    :stdio
    |> IO.read(:all)
    |> solve()
    |> IO.puts()
  end

  defp split_lines(lines) do
    lines
    |> String.trim()
    |> String.split("\n")
  end

  defn calc(tensor_500, tensor_100, tensor_50, total) do
    Nx.multiply(tensor_500, 500)
    |> Nx.new_axis(0)
    |> Nx.transpose()
    |> Nx.add(Nx.multiply(tensor_100, 100))
    |> Nx.new_axis(0)
    |> Nx.transpose()
    |> Nx.add(Nx.multiply(tensor_50, 50))
    |> Nx.equal(total)
    |> Nx.sum()
  end

  def solve(input) do
    [num_500, num_100, num_50, total] =
      input
      |> split_lines()
      |> Enum.map(&String.to_integer/1)

    tensor_500 = Nx.iota({num_500 + 1})
    tensor_100 = Nx.iota({num_100 + 1})
    tensor_50 = Nx.iota({num_50 + 1})

    calc(tensor_500, tensor_100, tensor_50, total)
    |> Nx.to_number()
  end
end

実行時間は 1718 msec でした

まとめ

Nx を使って組み合わせの問題を解きました

動的計画法よりも、やっていること自体は直観的で単純です

行列演算らしくできたので満足です