AtCoder ABC331をElixirで復習してみよう

問題A

caseなどの条件文でも書けるとは思いますが、関数化したほうがよさそうな気がします。
next_dayとnext_monthの関数に分けて実装してみました。

defmodule Main do
    import Bitwise
    def next_token(acc \\ "") do
        case IO.getn(:stdio, "", 1) do
          " " -> acc
          "\n" -> acc
          x -> next_token(acc <> x)
        end
    end
    def input(), do: IO.read(:line) |> String.trim()
    def ii(), do: next_token() |> String.to_integer()
    def li(), do: input() |> String.split(" ") |> Enum.map(&String.to_integer/1)

    def next_day(y,m,d, end_of_day, end_of_month) when d < end_of_day, do: {y,m,d+1}
    def next_day(y,m,d, end_of_day, end_of_month) when d == end_of_day, do: next_month(y,m, end_of_month)

    def next_month(y,m, end_of_month) when m < end_of_month, do: {y,m+1,1}
    def next_month(y,m, end_of_month) when m == end_of_month, do: {y+1,1,1}

    def solve(m, d, y, end_of_month, end_of_day) do
        {y,m,d} = next_day(y,m,d, end_of_day, end_of_month)
        IO.puts("#{y} #{m} #{d}")
    end

    def main() do
        end_of_month = ii()
        end_of_day = ii()
        y = ii()
        m = ii()
        d = ii()
        solve(m, d, y, end_of_month, end_of_day)
    end

end

問題B

方針

• Nが小さいので、各サイズのパック数を0~100について試しても十分間に合う。（100までする必要もないけど、気にしない）
• 全探索する

Elixirで全探索をどう書くかですが、for文で書いてみます。
１つのfor文でi,j,kの3重ループが書けるのが気に入ってるので、使ってみました。
この例の様に、ループの最初から最後まで通して参照する値があるときには、reduceを使うのが定石です。
こんな時、for文のreduceオプションは重宝します。

    def solve(n, s, m, l) do
        for i <- 0..n, j <- 0..n, k <- 0..n, reduce: 10**18 do
            acc ->
                cond do
                    i*6 + j*8 + k*12 >= n -> min(acc, s*i + m*j + l*k)
                    true -> acc
                end
        end
        |> IO.puts()
    end

    def main() do
        n = ii()
        s = ii()
        m = ii()
        l = ii()
        solve(n, s, m, l)
    end

問題C

方針

• Aの値をキー、出現回数を値とするmapを作る
• キーの値の小さい順にならべる
• キーの値×出現回数の累積和を求める

    def solve(_n, a) do

        cnt = counter(a)
        # IO.inspect(cnt)
        sorted_key = cnt |> Map.keys() |> Enum.sort()

        acc_sum_list=
        sorted_key
        |> Enum.map(fn x -> cnt[x]*x end)
        |> Enum.reduce([], fn x, acc -> [x+get_head(acc)|acc] end)
        |> Enum.reverse()

        sum_all = List.last(acc_sum_list)
        table =
            for {k, acc_sum} <- Enum.zip(sorted_key, acc_sum_list) , into: %{} do
                {k, sum_all - acc_sum}
            end

        a
        |> Enum.map(fn x -> table[x] end)
        |> Enum.join(" ")
        |> IO.puts()
    end

    def main() do
        n = ii()
        a = li()  # type: "List[int]"
        solve(n, a)
    end

Elixirで書くの難しそうな気がしたけど、意外と簡素に書けました。
累積和は関数があった気もするんですが、reduceで簡単に書けました。

参考　Pythonのコード

N = II()
A = LII()

cnt = Counter(A)
key = list(cnt.keys())
key.sort()
sum_list = [cnt[k]*k for k in key]
acc_sum_list = list(accumulate(sum_list))
table = {}
for i,k in enumerate(key):
    table[k] = acc_sum_list[-1]-acc_sum_list[i]

ans = []
for i in range(N):
    k = A[i]
    ans.append((table[k]))

print(*ans)

問題D

まだ解いてないです。後で解きます。
二次元の累積和でもとめられそうな問題

問題E

方針

• 主菜を１つづつ試す
• 副菜を大きい順に選んで、主菜、副菜が禁止リストにないものを見つけて最大値を更新する。

副菜の繰り返しは、禁止ではない組み合わせが見つかったら、中断する必要があります。reduce_whileを使いました。
再帰関数で求めてもよいとおもいます。
Pythonの場合は、break文でループを抜ける記述でできますが、Elixirの場合慣れが必要ですね。

PythonではACのコード作れたんですが、Elixirでは、TLEになってしまいました。
入力の量が多いので、ちょっと入力の処理を改善しないといけないかもしれません。

    def solve(n, m, l, a, b, cd) do

        cd_set =
        for [c,d] <- cd, reduce: MapSet.new() do
            acc -> MapSet.put(acc, {c-1,d-1})
        end

        sa = create_sorted_value_index_list(a)
        sb = create_sorted_value_index_list(b)

        # IO.inspect(sa)
        # IO.inspect(sb)
        # IO.inspect(cd_set)

        for {av,ai} <- sa, reduce: 0 do
            acc ->
                Enum.reduce_while(sb, acc, fn {bv,bi}, acc ->
                    cond do
                        MapSet.member?(cd_set, {ai,bi}) -> {:cont, acc}
                        true -> {:halt, max(acc,av+bv)}
                    end
                end)
        end
        |> IO.puts()
    end

    def main() do
        n = ii()
        m = ii()
        l = ii()
        a = li()
        b = li()
        cd =
        if l >0 do
            for _ <- 1..l, do: li()
        else
            []
        end
        solve(n, m, l, a, b, cd)
    end

参考　Pythonのコード

N,M,L = TII()
A = LII()
B = LII()
cd = [LII() for _ in range(L)]

SA = [(v,i) for i,v in enumerate(A)]
SB = [(v,i) for i,v in enumerate(B)]

SA.sort(reverse=True)
SB.sort(reverse=True)

set_cd = set()
for c,d in cd:
    c -= 1
    d -= 1
    set_cd.add((c,d))

# print(SA)
# print(SB)
# print(set_cd)
ans = 0
for a in SA:
    for b in SB:
        if (a[1],b[1]) not in set_cd:
            ans = max(ans,a[0]+b[0])
            break
print(ans)