ABC081B - Shift only

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

問題

回答3

:math.pow(2, 29)

powers =
  Nx.pow(2, Nx.iota({30}))

[8, 12, 40]
|> Nx.tensor()
|> Nx.new_axis(0)
|> Nx.transpose()
|> Nx.remainder(powers)
|> Nx.reverse(axes: [1])
|> Nx.argmin(axis: 1)
|> then(&Nx.subtract(29, &1))
|> Nx.reduce_min()
|> dbg()

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

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

  defn get_powers() do
    2 ** Nx.iota({30})
  end

  defn calc(input_tensor, powers) do
    input_tensor
    |> Nx.new_axis(0)
    |> Nx.transpose()
    |> Nx.remainder(powers)
    |> Nx.reverse(axes: [1])
    |> Nx.argmin(axis: 1)
    |> then(&Nx.subtract(29, &1))
    |> Nx.reduce_min()
  end

  def solve(input) do
    nums =
      input
      |> split_lines()
      |> Enum.at(1)
      |> split_words()
      |> Enum.map(&String.to_integer/1)
      |> Nx.tensor()

    powers = get_powers()

    nums
    |> calc(powers)
    |> Nx.to_number()
  end
end

"""
3
8 12 40
"""
|> Main.solve()
|> then(&(&1 == 2))

"""
4
5 6 8 10
"""
|> Main.solve()
|> then(&(&1 == 0))

"""
6
382253568 723152896 37802240 379425024 404894720 471526144
"""
|> Main.solve()
|> then(&(&1 == 8))