ABC130D - Enough Array

問題

回答（TLE）

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)

    renew_table([], :sum)
    renew_table([], :false_combination)

    false_combinations =
      for i <- 0..(n - 1),
          j <- i..(n - 1),
          a_j = lookup(:a, j),
          pre_sum = lookup(:sum, {i, j - 1}),
          sum = pre_sum + a_j,
          :ets.insert(:sum, {{i, j}, sum}),
          sum < k do
        1
      end
      |> length()

    all_combination = div(n * (n + 1), 2)

    all_combination - false_combinations
  end
end

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

:false_combination
|> :ets.tab2list()
|> Enum.sort_by(&elem(&1, 0))

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

:false_combination
|> :ets.tab2list()
|> Enum.sort_by(&elem(&1, 0))

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