ABC083B - Some Sums

Mix.install(
  [
    {:nx, "~> 0.9"},
    {:exla, "~> 0.9"},
    {:kino, "~> 0.14"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ]
  ]
)

問題

回答3

nums = Nx.iota({1235})

digit = Nx.pow(10, Nx.iota({5})) |> Nx.reverse()

sums =
  nums
  |> Nx.new_axis(0)
  |> Nx.transpose()
  |> Nx.quotient(digit)
  |> Nx.remainder(10)
  |> Nx.slice([995, 0], [10, 5])

defmodule Main do
  import Nx.Defn

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

  defp split_words(words) do
    String.split(words, " ")
  end

  defn calc(nums, digit, a, b) do
    nums
    |> Nx.new_axis(0)
    |> Nx.transpose()
    |> Nx.quotient(digit)
    |> Nx.remainder(10)
    |> Nx.sum(axes: [1])
    |> then(&Nx.logical_and(Nx.greater_equal(&1, a), Nx.less_equal(&1, b)))
    |> Nx.multiply(nums)
    |> Nx.sum()
  end

  def solve(input) do
    [n, a, b] =
      input
      |> String.trim()
      |> split_words()
      |> Enum.map(&String.to_integer/1)

    nums = Nx.iota({n + 1})
    digit = Nx.pow(10, Nx.iota({5}))

    calc(nums, digit, a, b)
    |> Nx.to_number()
  end
end

"""
20 2 5
"""
|> Main.solve()
|> then(&(&1 == 84))

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

"""
100 4 16
"""
|> Main.solve()
|> then(&(&1 == 4554))