[AtCoder Nx] ABC087B - Coins [Elixir Livebook]

## はじめに

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

https://qiita.com/RyoWakabayashi/items/206a19dc53e334ca7b60

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

https://qiita.com/RyoWakabayashi/items/69d8c5318bf2cef8b74e

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/atcoder/abc087/abc087_b/submit_2.livemd

## 問題

https://atcoder.jp/contests/abs/tasks/abc087_b

## 入出力例

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

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

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

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

## Nx なしの回答

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

```elixir
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] の全ての和の組み合わせは以下のように計算します


```math
\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 としたとき、使った枚数毎の金額を表す行列は以下のコードで表せます

:::note info
`Nx.iota(x)` = [0, 1, 2, ..., x - 1]
:::

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

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

```elixir
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 を使って組み合わせの問題を解きました

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

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