ABC242のA, B, C問題 を Elixir で解く

これに続いて、Elixir 超初心者が AtCoder に挑戦してみました。

クソコードを御指導下さい。

A問題: T-shirt

defmodule Main do
  def main do
    IO.read(:line) |> integer_list() |> do_solve() |> IO.puts()
  end
  
  defp integer_list(str) do
    str
    |> String.trim()
    |> String.split(" ")
    |> Enum.map(&String.to_integer/1)
  end
  
  defp do_solve([a, _, _, x]) when 1 <= x and x <= a, do: 1.0
  defp do_solve([a, b, c, x]) when a < x and x <= b, do: c / (b - a)
  defp do_solve(_), do: 0
end

ガード節の中では&& ||が使えず、and orを使うというのがハマリポイントでした。それから Elixir では整数どうしの/でも結果がFloatになるんですね。

B問題: Minimize Ordering

defmodule Main do
  def main do
    IO.read(:line)
    |> String.trim()
    |> String.to_charlist()
    |> Enum.sort()
    |> List.to_string()
    |> IO.puts()
  end
end

これは文字列をコードポイントに直してソートして、また文字列に直すと。

C問題: 1111gal password

これは動的計画法なんですが、はて、関数型言語で動的計画法とは…となりました。ぐぐってみると、reduceを使うと。Ruby ならinjectで（reduceもありますが）、Elixir ならEnum.reduceですね。なるほど、と。

で、リストを使って実装したところ、

defmodule Main do
  def main do
    IO.read(:line)
    |> String.trim() |> String.to_integer()
    |> do_solve([1, 1, 1, 1, 1, 1, 1, 1, 1])
    |> IO.puts()
  end
  
  defp do_solve(n, list) do
    1..n - 1
    |> Enum.reduce(list, &get_next_list(&1, &2))
    |> Enum.sum() |> r()
  end
  
  defp get_next_list(_, list) do
    0..8 |> Enum.map(&dp(&1, list))
  end
  
  defp dp(0, enum), do: Enum.at(enum, 0) + Enum.at(enum, 1) |> r()
  defp dp(8, enum), do: Enum.at(enum, 7) + Enum.at(enum, 8) |> r()
  defp dp(i, enum) do
    Enum.at(enum, i - 1) + Enum.at(enum, i) + Enum.at(enum, i + 1)
    |> r()
  end
  
  defp r(n), do: rem(n, 998244353)
end

TLEをくらいました。どうも、Enum.atが遅いのではと考え、マップを使って再挑戦。マップは Ruby だとハッシュって感じですね。いわゆる連想配列。

defmodule Main do
  def main do
    IO.read(:line) |> String.trim() |> String.to_integer()
    |> do_solve(%{1=>1, 2=>1, 3=>1, 4=>1, 5=>1, 6=>1, 7=>1, 8=>1, 9=>1})
    |> IO.puts()
  end
  
  defp do_solve(n, map) do
    1..n - 1
    |> Enum.reduce(map, &get_next_map(&1, &2))
    |> Map.values()
    |> Enum.sum() |> r()
  end
  
  defp get_next_map(_, map) do
    1..9
    |> Enum.reduce(%{}, fn i, acc -> Map.put(acc, i, dp(i, map)) end)
  end
  
  defp dp(1, map), do: map[1] + map[2] |> r()
  defp dp(9, map), do: map[8] + map[9] |> r()
  defp dp(i, map), do: map[i - 1] + map[i] + map[i + 1] |> r()
  
  defp r(n), do: rem(n, 998244353)
end

なんとか AC しました。1652ms ですけれどね。なかなかむずかしい。

最後に

もっといい書き方があったら教えて下さいね。