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

AtCoder Beginner Contest 330をElixirで解いてみます

問題A

リストAの中で、値がL以上のものが何個あるかカウントする問題。
Enumで解くのにピッタリ

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 main() do
        n = ii()
        l = ii()
        a = li()
        ans = Enum.count(a, fn x -> x >=l end)
        IO.puts(ans)
    end

end

問題B

この問題、Bなので簡単なはずなのに、問題を理解できず、プログラミング以前で悩みました。
10分は浪費しました。落ち着いて考えると、AiとL,Rの値からXiが決まります。

条件	Xiの値
AiがLより小さい	L
AiがL~Rの間	Ai
AiがRより大きい	R

この法則がわかれば、Elixirで書けます

方針

• 条件（この場合、L,R,Ai)の値を与えるとXiを返す関数(get_x)を作る
• Aのリスト全体にこの関数を適用してXを求める

get_xは、パターンマッチで書けそう。

リストの数値をIO.putsで表示する方法を忘れてたので、過去の回答から探してきました。

    def get_x(l,_r,a) when a<l, do: l
    def get_x(_l,r,a) when a>r, do: r
    def get_x(_l,_r,a), do: a

    def main() do
        n = ii()
        l = ii()
        r = ii()
        a = li()

        Enum.map(a, fn a -> get_x(l,r,a) end)
        |> Enum.join(" ")
        |> IO.puts()
    end

問題C

Xでは、尺取り法での解説をみかけました。私は、Pythonでは、二分探査で解いたので、二分探査で解いてみます。尺取り方をElixirで書くのって、難しそう。

方針

• xを与えたとき、xx+yy<Dを満たす最小のyを二分探査で求める関数get_y(x,d)を作成する
• xを0からsqrt(D)の範囲について、min(f(y),f(y+1))の最小値を求める
  y = get_y(x,d)
  f(y) = abs(xx+yy-D)

xの全探索は、for文のreduceで書いてみます

    def get_y(x,d,func,l,r) when r-l==1 do
        l
    end

    def get_y(x,d,func,l,r) do
        m = div((l+r),2)
        if func.(m) do
            get_y(x,d,func,m,r)
        else
            get_y(x,d,func,l,m)
        end
    end

    def get_y(x,d) do
        func = fn y -> x*x + y*y < d end
        l = 0
        r = floor(:math.sqrt(d)+1)
        get_y(x,d,func,l,r)
    end

    def xyd(x,y,d) do
        abs(x*x + y*y - d)
    end

    def main() do
        d = ii()
        for x <- 0..floor(:math.sqrt(d)), reduce: D do
            acc ->
                y = get_y(x,d)
                a = min(xyd(x,y,d), xyd(x,y+1,d))
                min(a, acc)
        end
        |> IO.puts()
    end

二分探査をElixirで初めて書いてみたんですが、再帰との相性がいいですね
TLEにならないか心配だったけど、1394msでACできました。

問題D

D問題になってくると、計算量の適切なアルゴリズムを思いつくのが一苦労です。
入力例を見ながら規則性がわかったところでプログラムにしてみます。

方針

• cnt_h[i] = i行のoの数のリストを作成
• cnt_w[j] = j列のoの数のリストを作成
• s[i][j]=="o"のi,jについて (n_h[i] -1) * (n_w[j] -1)の合計を求める

方針の3つ目は、Pythonぽい書き方で書きましたが、実際次のようなPythonプログラムでACできます

ans = 0
for i in range(N):
    for j in range(N):
        if S[i][j] != "o":
            continue
        ans += (cnt_h[i] -1) * (cnt_w[j] -1)

Elixirのリストは、要素の参照がO(N)かかるので、使えません。tupleに変換してelemで要素を参照する方法もありますが、記述が結構大変。
次のようなプログラムにしてみました。

    def count_char(line,c) do
        Enum.count(line, fn x -> x == c end)
    end

    def count_o(s) do
      Enum.map(s, fn line -> count_char(line, ?o) end)
    end

    def transpose(s) do
        Enum.zip(s) |> Enum.map(&Tuple.to_list/1)
    end

    def main() do
        n = ii()
        s = for _ <- 1..n, do: input()|> String.to_charlist()
        cnt_h = count_o(s)
        cnt_v = count_o(transpose(s))

        flat_s = List.flatten(s)
        flat_cnt_v = List.duplicate(cnt_v, n) |> List.flatten()
        flat_cnt_h = Enum.flat_map(cnt_h, fn x -> List.duplicate(x, n) end)
        for {c, h, v} <- Enum.zip([flat_s, flat_cnt_h, flat_cnt_v]), reduce: 0 do
            acc -> if c == ?o, do: acc + (h-1) * (v-1), else: acc
        end
        |> IO.puts()
    end

残念ながらこのコードではTLEでした。
長さが20002000=410^6のリストを扱うことになるので、結構厳しそう。
Nxを使うとかStreamを使うとかが必要かもしれません。
ACできるコードは、また別途検討してみます。

問題E

Pythonでは、SortedSetを使って解きました。
Elixirではまだ未開拓な問題です。できるようになったら記事にしようとおもいます。