ABC130D - Enough Array

問題

回答

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

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

  defp renew_table(objects, table_name) do
    if :ets.info(table_name) != :undefined do
      :ets.delete(table_name)
    end

    :ets.new(table_name, [:set, :protected, :named_table])
    :ets.insert(table_name, objects)
  end

  defp lookup(table_name, key) do
    case :ets.lookup(table_name, key) do
      [{_, value}] -> value
      _ -> 0
    end
  end

  def solve(input) do
    [[n, k], a] =
      input
      |> split_lines()
      |> Enum.map(fn line ->
        line
        |> split_words()
        |> Enum.map(&String.to_integer/1)
      end)

    a
    |> Enum.with_index()
    |> Enum.map(fn {a_i, i} -> {i, a_i} end)
    |> renew_table(:a)

    {combination_sum, _, _} =
      0..(n - 1)
      |> Enum.reduce_while({0, 0, -1}, fn i, {combination, acc_i, pre_head_j} ->
        # i: 部分配列の左端
        # acc_i: 部分配列の合計
        # pre_head_j: 部分配列の右端
        cond do
          acc_i < k &amp;&amp; pre_head_j == n - 1 ->
            # 最後まで計算して、もう合計が k 以上になることがない場合、終了する
            {:halt, {combination, 0, 0}}

          acc_i < k ->
            # k 以上になるまで部分配列の右端をずらす
            # i から pre_head_j までの部分配列の合計は acc_i に入っているので、それより後だけ計算する
            init_j = pre_head_j + 1

            {head_j, sum_j, new_combination} =
              init_j..(n - 1)
              |> Enum.reduce_while({init_j, acc_i, 0}, fn j, {_, acc_j, _} ->
                sum_j = acc_j + lookup(:a, j)

                if sum_j < k do
                  {:cont, {j, sum_j, 0}}
                else
                  # k 以上になったとき、 i から j 以降の合計は全て k 以上になる
                  {:halt, {j, sum_j, n - j}}
                end
              end)

            new_head_j = if head_j > pre_head_j, do: head_j, else: pre_head_j
            {:cont, {combination + new_combination, sum_j - lookup(:a, i), new_head_j}}

          true ->
            # k 以下になるまで部分配列の左端をずらす
            # k 以下になるまで、 i から pre_head_j 以降の合計は全て k 以上になる
            new_combination = n - pre_head_j
            {:cont, {combination + new_combination, acc_i - lookup(:a, i), pre_head_j}}
        end
      end)

    combination_sum
  end
end

"""
4 10
6 1 2 7
"""
|> Main.solve()
|> then(&(&1 == 2))

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

"""
10 53462
103 35322 232 342 21099 90000 18843 9010 35221 19352
"""
|> Main.solve()
|> then(&(&1 == 36))