AtCoder に登録したら解くべき精選過去問 10 問を Prolog で提出してみた

概要

そう言えば AtCoder Language Test にPrologが追加されていましたね。
「精選過去問 10 問」をPrologで解くということはすでに 「先行研究」AtCoder に登録したら解くべき精選過去問 10 問を Prolog で解いてみた [1] がありますが実行環境の違いや制約もあるので提出が通るかどうか試してみました。
コードをそのまま提出するだけというのも芸がないので、先行研究¹を参考にしつつ、極力自分で解いて一言コメントをつけてみようと思います。

0. practice contest A - Welcome to AtCoder

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10481318 ²

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).
get_number(Number) :-
  get_string(String),
  number_string(Number, String).

main :-
  get_number(A),
  get_number(B),
  get_number(C),
  get_string(S),
  Sum is A + B + C,
  write(Sum),
  write(' '),
  write(S).

実行時間 102 ms
メモリ 8152 KB

先行研究¹とほぼ同じになってしまいました。ジャッジの際、自動で main を呼んでくれるようなので:-main.はあってもなくても動きます。(main が無いとCEになる。)

代入(というか変数束縛)はisを使っていますが、ここが=だとA + B + Cという式ごとUnifyされ、「1+2+3 test」のような出力になってしまいます。isだと数値演算をしてから代入されます。

1. ABC 086 A - Product

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10481612

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).
get_number(Number) :-
  get_string(String),
  number_string(Number, String).

main :-
  get_number(A),
  get_number(B),
  A * B mod 2 =:= 0 ->
    write('Even');
    write('Odd').

実行時間 99 ms
メモリ 8052 KB

if文は _if -> _true; _false.のように書けます。
等値比較は=:=です。isと同様、演算結果を比較します。

2. ABC 081 A - Placing Marbles

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10541880

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

count1([], 0).
count1(['1'| Rest], Ans):-
  count1(Rest, AnsR),
  Ans is AnsR + 1.
count1(['0'| Rest], Ans):-
  count1(Rest, Ans).

main :-
  get_string(S),
  string_chars(S, Chars),
  count1(Chars, Ans),
  write(Ans).

実行時間 59 ms
メモリ 8116 KB

文字列を文字のリストに変換し、先頭から'1'が含まれている数をカウントします。

3. ABC 081 B - Shift Only

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10482609

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

get_numbers(0, []).
get_numbers(Size, [Number1 | List0]):-
  Size0 is Size - 1,
  get_numbers(Size0, List0),
  get_number(Number1).


even(X) :- X mod 2 =:= 0.

half(X, Y) :- Y is X // 2.

shift_time(N, List):-
  maplist(even, List) ->
    maplist(half, List, List0),
    shift_time(N0, List0),
    N is N0 + 1;
    N = 0.

main :-
  get_number(N),
  get_numbers(N, List),
  shift_time(Ans, List),
  write(Ans).

実行時間 107 ms
メモリ 8284 KB

maplist/3は第1引数に2項の述語を取り、第2引数、第3引数の各要素を適用してUnifyする高階述語です。hoge(X,Y)をhoge:X→Y という関数と見て写像してやるような操作(PythonやScalaのmapなどのような)に使えますが、hoge/2が双方向で使えればこれも双方向に使えるというのがprologらしいところですね。
plus/3はplus(A,B,C)でA+B=Cとなる数字をマッチしますが、plus(1)として「1を足す」という述語を作ることも可能です。これも関数型っぽいですね。

maplist/2は1項の述語を取って、第2引数の全要素についてtrueならtrueとなる述語で、これもPythonのallやScalaのforallなどに似ています。

4. ABC 087 B - Coins

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10482821

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

ok(A, B, C, P, Q, R, X):-
  between(0, A, P),
  between(0, B, Q),
  between(0, C, R),
  X =:= 500 * P + 100 * Q + 50 * R.

main :-
  get_number(A),
  get_number(B),
  get_number(C),
  get_number(X),
  findall(_, ok(A, B, C, P, Q, R, X), Results),
  length(Results, Ans),
  write(Ans).

実行時間 115 ms
メモリ 8176 KB

特に別解が思いつかなかったので先行研究¹とほぼ同じになりました。

findall/3はその名の通り解をすべて見つける述語で、findall(X, hoge(X, Fuga, Piyo), Res)のようにすると、hogeが成立するXを全て見つけてResに格納します。
今回は数だけ分かればいいので値は_で捨てています。

5. ABC 083 B - Some Sums

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10502174

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

digit_sum(0, 0):- !.
digit_sum(Number, Dsum):-
  divmod(Number, 10, Q, R),
  digit_sum(Q, Dsum0),
  Dsum is R + Dsum0.

ok(N, A, B, X):-
  between(1, N, X),
  digit_sum(X, Dsum),
  between(A, B, Dsum).

main:-
  get_number(N),
  get_number(A),
  get_number(B),
  findall(X, ok(N, A, B, X), Results),
  sum_list(Results, Sum),
  write(Sum).

実行時間 71 ms
メモリ 8580 KB

これも別解が思いつかなかったので先行研究¹とほぼ同じです。

divmod/4はdivmod(N,P,Q,R)でQ=N//P(商)、R=N%P(剰余)を同時に計算してくれます。

6.ABC 088 B - Card Game for Two

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10502349

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

get_numbers(0, []).
get_numbers(N, [X|Xs]) :-
  get_number(X),
  N1 is N - 1,
  get_numbers(N1, Xs).

score(_, [], 0).
score(0, [A|Rest], Score):-
  score(1, Rest, Score1),
  Score is Score1 + A.
score(1, [A|Rest], Score):-
  score(0, Rest, Score1),
  Score is Score1 - A.

main:-
  get_number(N),
  get_numbers(N, As),
  sort(0, @>=, As, SortedAs),
  score(0, SortedAs, Score),
  write(Score).

実行時間 62 ms
メモリ 8260 KB

Aliceの番とBobの番、で交互に再帰してみました。

sort/2は重複を捨てた昇順ソートになります。重複を保つにはsort/4で@<=や@>=で比較演算子を指定する必要があります。

7. ABC 085 B - Kagami Mochi

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10502453

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

get_numbers(0, []).
get_numbers(N, [X|Xs]) :-
  get_number(X),
  N1 is N - 1,
  get_numbers(N1, Xs).

main:-
  get_number(N),
  get_numbers(N, Ds),
  sort(Ds, SortedDs),
  length(SortedDs, Ans),
  write(Ans).

実行時間 59 ms
メモリ 8108 KB

これも別解が思いつかなかったので先行研究¹とほぼ同じです。

sort/2が重複を除くのを利用して個数を数えます。

8. ABC 085 C - Otoshidama

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10540229

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

ok(N, Y, P, Q, R):-
  between(0, N, P),
  Rest is N - P,
  between(0, Rest, Q),
  R is N - P - Q,
  Y =:= 10000 * P + 5000 * Q + 1000 * R, !.

main :-
  get_number(N),
  get_number(Y),
  ok(N, Y, P, Q, R)->
    write(P), write(" "), write(Q), write(" "), write(R);
    write("-1 -1 -1").

実行時間 438 ms
メモリ 8132 KB

P,QからRは確定するのを利用して$O(N^2)$にして計算量を落とします。
答えが一個見つかれば良いので一応最後にカット!を入れましたが、無くても動くのであまりよくわかっていません。

愚直にP,Q,Rを全て探索する$O(N^3)$解法では無事(?)TLEになりました。
https://atcoder.jp/contests/language-test-202001/submissions/10539698

コード内容

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

ok(N, Y, P, Q, R):-
  between(0, N, P),
  between(0, N, Q),
  between(0, N, R),
  plus(P, Q, PQ), plus(PQ, R, N),
  Y =:= 10000 * P + 5000 * Q + 1000 * R.

main :-
  get_number(N),
  get_number(Y),
  findall((P, Q, R), ok(N, Y, P, Q, R), Results),
  [(A, B, C)|_] = Results->
    write(A), write(" "), write(B), write(" "), write(C);
    write("-1 -1 -1").

実行時間 2239 ms
メモリ 8100 KB

9. ABC 049 C - 白昼夢

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10539927

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

ok([]).
ok([m, a, e, r, d | Rest]) :- ok(Rest).
ok([r, e, m, a, e, r, d | Rest]) :- ok(Rest).
ok([e, s, a, r, e | Rest]) :- ok(Rest).
ok([r, e, s, a, r ,e | Rest]) :- ok(Rest).

main :-
  get_string(S),
  string_chars(S, Chars),
  reverse(Chars, RevChars),
  ok(RevChars) ->
    write('YES');
    write('NO').

実行時間 79 ms
メモリ 15732 KB

逆からマッチしていく想定解通りの実装をしてみましたが、実は頭からマッチングするように書いても(＝先行研究¹の解法)通ります。
https://atcoder.jp/contests/language-test-202001/submissions/10539942

コード内容

% https://qiita.com/n4o847/items/2040067b5388481eae88#6-abc088b---card-game-for-two より

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

satisfy([]).
satisfy([d, r, e, a, m | Rest]) :- satisfy(Rest).
satisfy([d, r, e, a, m, e, r | Rest]) :- satisfy(Rest).
satisfy([e, r, a, s, e | Rest]) :- satisfy(Rest).
satisfy([e, r, a, s, e, r | Rest]) :- satisfy(Rest).

main :-
  get_string(S),
  string_chars(S, Chars),
  satisfy(Chars) ->
    write('YES');
    write('NO').

実行時間 63 ms
メモリ 11976 KB

しかも想定解より若干速い。
Prologすごいですね…… 計算量をどう見積もれば良いのか全然わかりません。

10. ABC 086 C - Traveling

提出コード
https://atcoder.jp/contests/language-test-202001/submissions/10540177

get_string(String) :-
  current_input(Input),
  read_string(Input, ' \n', '', _, String).

get_number(Number) :-
  get_string(String),
  number_string(Number, String).

solve(0, _, _, _).
solve(N, T0, X0, Y0) :-
  get_number(T),
  get_number(X),
  get_number(Y),
  Dt = T - T0,
  L is abs(X - X0) + abs(Y - Y0),
  Dt >= L,
  Dt mod 2 =:= L mod 2,
  N1 is N - 1,
  solve(N1, T, X, Y).

main:-
  get_number(N),
  solve(N, 0, 0, 0)->
    write("Yes");
    write("No").

実行時間 231 ms
メモリ 8164 KB

これも別解が思いつかなかったので先行研究¹とほぼ同じです。

終わりに

1年前にPrologが話題になった時にも一度やろうとしていたのですが当時は高階述語あたりの使い方がよくわからなくて挫折していました。
関数型をちょっと触ったのがいい経験になったのか結構分かるようになりました。関数型の考え方とは親和性が高いと思います。

AtCoderで採用されているSWI-PrologのDocumentはこちら https://www.swi-prolog.org/ で、述語のAPIなどが調べられます。
意外な述語が双方向定義だったりして面白いです。
例えばplus/3のDocumentを見てみると

plus(?Int1, ?Int2, ?Int3)
True if Int3 = Int1 + Int2. At least two of the three arguments must be instantiated to integers.

と書いてあり、3つのうち2つが確定していればもう1つはどれでも求められ、plus(1, X, 3)で3-1が計算できたりします。

実際に提出してみると思ったよりも高速で驚きました。そのへんのスクリプト言語よりかは速いんじゃないでしょうか。
特に「白昼夢」は頭から探索するよう書いてもちゃんと高速に処理できていてすごかったですね。
このような「Prologが強い」問題もあると思うので正式採用されると面白そうだなあと思っています。

---

1. AtCoder に登録したら解くべき精選過去問 10 問を Prolog で解いてみた @n4o847 さんより ↩

2. 提出したコードには get_string(T)(l13) というのが入っていますが、EOLを読んだらどうなるのか実験をしていたのを消し忘れただけなので気にしないでください。(EOLだとread_string/5の第4引数が-1になり(ここでは_で捨てられている)、Stringには""が入ります。) ↩