LoginSignup
5
3

More than 3 years have passed since last update.

@selpo

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

Posted at

概要

drken さんの 過去問精選 10 問 を、Wolfram で解いてみました。

Wolfram とは？

Wolfram 言語とは、Mathematica で用いられているマルチパラダイムな言語です。Wolfram|Alpha の中でも動いていたりします。

基本的には項書き換えシステムで、関数型っぽい雰囲気もあり、手続き型であるかのようにも使えます。Lisp の方言 と言われることもあります。クセが強い言語です。

ところで Mathematica は有償です 1。学生版でも これくらい します。そのため、基本的には、大学などからアカウントをもらえる人くらいしか使いません。Wolfram 言語を学んでみようと思っても、気軽に試すことができませんでした。

しかし、先日 Wolfram Engine が開発者向けに無償で公開されました！これにより、Wolfram 言語で書かれたスクリプトを誰でも実行してみることができます 2。

ということで、手元でも実行してみたいという方は、ここ に従って Wolfram Engine をインストールしておいてください。wolframscript コマンドが使えるようになっているはずです (1 回目は認証が必要です)。

第 0 問: practice A - Welcome to AtCoder

まずは入出力の練習をしましょう。

0.wls
main[] := Module[{lines, lineno = 1, getLine, a, b, c, s},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    a = ToExpression @ getLine[];
    {b, c} = ToExpression /@ StringSplit @ getLine[];
    s = getLine[];
    Print @ StringRiffle[{a + b + c, s}];
]
main[]

実行方法

プログラムの解説をする前に、実行方法を説明します。

これを 0.wls に保存し、入力ファイルを in0-1.txt に用意しておきましょう。

in0-1.txt
1
2 3
test

実行するには、wolframscript -file 0.wls < in0-1.txt とします。すると、6 test と表示されるはずです。

ちなみに、shebang をつけることもできます。0.wls の最初の行に #!/usr/bin/env wolframscript をつけて chmod +x 0.wls としておくと、./0.wls < in0-1.txt でも実行できます。

解説

最初の 8 行 (main[] := Module[...]) は, main[] という関数を定義しています。別に main という特別な関数を定義する必要があるわけではないです。main の中身をべた書きしても OK です。

ちなみに、Mathematica では関数呼び出しは () ではなく [] を使います。() は演算子の結合順の指定でしか使いません。配列などのアクセスは [[]] を使います。

9 行目の main[] で実際に main 関数を呼び出しています。正確には、main[] という項を評価しようとすると、main[] := Module[...] という評価規則が与えられていたので、main[] を Module[...] に書き換えてそれを評価しています。これ以上書き換えができない形までもっていくことが、計算の基本になります。

Module 関数は、局所変数を使えるようにしてくれます。Module[{x,y,...}, expr] とすると、expr の評価中に x,y,... を局所変数として使うことができます。実際の実行時には、 x$123,y$123,... のようにユニークな変数が生成されます。初期値を設定することもできて、lineno=1 を与えています。

expr を見てみると、何やら ; で式が区切られていますね。これは CompoundExpression といって、expr1;expr2;... は CompoundExpression[expr1, expr2, ...] のシンタックスシュガーで、expr1, expr2, … を順番に評価して、最後に評価された式の値を返します 3。要は、C や C++ のカンマ演算子です。

Wolfram 言語はシンタックスシュガーが豊富で、0.wls に出てきたものだけでも、

  • lhs := rhs は SetDelayed[lhs, rhs] → SetDelayed
  • lhs = rhs は Set[lhs, rhs] → Set
  • {e1, e2, ...} は List[e1, e2, ...] → List
  • expr[[i]] は Part[expr, i] → Part
  • x++ は Increment[x] → Increment
  • f @ x は f[x]
  • f /@ expr は Map[f, expr] -> Map
  • x + y + ... は Plus[x, y, ...] -> Plus

これだけあります。この先でもさらに出てきます。逆に、シンタックスシュガーを全く用いないと [] とシンボルと、といくつかのプリミティブ (数値、文字列) しか存在しません 4。すべては式 なのです！！

話がそれました。解説に戻りましょう。
lines = StringSplit[$ScriptInputString, "\[NewLine]"] では、入力文字列 (今回は in0-1.txt の中身) を 改行文字 で split しています (何も指定しないと空白文字で split します)。結果は文字列の List で与えられ、

lines = {"1", "2 3", "test"};

のようになっていることでしょう。getLine[] は、呼び出すたびに lines の要素を順番に取ってくる関数として定義しました (Wolfram 言語では、添え字は 1 から始まることに注意。0 番目には式の 頭部 が入っている)。

a = ToExpression @ getLine[] で "1" を数値に変換しています。正確に言うと、ToExpression は文字列などを Wolfram 言語の式として 解釈します。要は eval です。

{b, c} = ToExpression /@ StringSplit @ getLine[] は、"2 3" を split し (→ {"2", "3"})、それぞれに ToExpression を適用します (→ {2, 3})。こういう風に List 同士の代入もできます(swap も {x, y} = {y, x} でできます)。

s = getLine[] はそのままですね。

あとは出力するだけです。

Print @ StringRiffle[{a + b + c, s}]

Print が出力関数です。空白区切りにしたいので、StringRiffle を用います。

第 1 問: ABC 086 A - Product

ようやく肩慣らし (?) が終わりました。長かったですが、ある程度 Wolfram 言語を理解してくれたと信じて、突き進みます。

2 つの整数の積が偶数かどうか判定します。

1.wls
main[] := Module[{lines, lineno = 1, getLine, a, b},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    {a, b} = ToExpression /@ StringSplit @ getLine[];
    Print @ If[EvenQ[a b], "Even", "Odd"];
]
main[]

If が必要になりますが、第 0 問が理解できていれば簡単ですね。

Wolfram 言語では、a b が a と b の積 (Times[a, b]) に解釈されます。数学で掛け算記号を省略する気分です。また、EvenQ、OddQ という偶奇を判定するだけの関数があるのでそれを用います。

If[cond, t, f] は、cond を評価して True になる場合は t を、 False になる場合は f を返します 5。返されない方は評価されません。

第 2 問: ABC 081 A - Placing Marbles

"101" のような文字列に含まれる 1 の個数を数えます。

2.wls
main[] := Module[{lines, lineno = 1, getLine, s},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    s = ToCharacterCode @ getLine[] - ToCharacterCode["0"][[1]];
    Print @ Total @ s;
]
main[]

s = ToCharacterCode @ getLine[] - ToCharacterCode["0"][[1]] は何をしているのでしょうか。ToCharacterCode は文字列を整数コードの List に変換します。 python の ord のようなもの ([ord(c) for c in s]) です。例えば、"101" が入力として与えられた場合、{49, 48, 49} です。ToCharacterCode["0"][[1]] は ord("0") に相当します。
すると、s = {49, 48, 49} - 48 が実行されることになります。List から定数を引く…? 実は、List と数値との四則演算は、各要素への演算になります。つまり、s = {49 - 48, 48 - 48, 49 - 48} です。
こうして、s = {1, 0, 1} になりました。そう、入力文字列を数字のリストに変換していたのでした。

あとは s の総和を求めて出力するだけです。Total という、総和を求める関数があるのでそれを用いましょう。

第 3 問: ABC 081 B - Shift Only

3.wls
main[] := Module[{lines, lineno = 1, getLine, N, A},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    N = ToExpression @ getLine[];
    A = ToExpression /@ StringSplit @ getLine[];
    Print @ Min[IntegerExponent[#, 2] & /@ A];
]
main[]

2 で最大何回割り切れるか求める必要があります。なんと、IntegerExponent がすでに用意されています。IntegerExponent[n, b] で「n が b で最大何回割り切れるか」を計算してくれます。
あとは Min でその最小値を答えるだけ…なのですが、なにやら怪しげな #、& がありますね。

body & は Function[body] のシンタックスシュガーで、純関数 を表します。要はラムダ式です。引数は # で参照できます6。
IntegerExponent[#, 2] & で「引数が 2 で最大何回割り切れるか」を返す関数になり、/@ を用いて A に Map しています。

ちなみに、全く同じことが IntegerExponent[A, 2] でもできます。Wolfram 言語の多くの関数は、引数が List のときには勝手に Map するように定義されているのです(今回は純関数を紹介するためにあえて用いませんでした)。

第 4 問: ABC 087 B - Coins

4.wls
main[] := Module[{lines, lineno = 1, getLine, A, B, C, X, cnt = 0, a, b, c},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    A = ToExpression @ getLine[];
    B = ToExpression @ getLine[];
    C = ToExpression @ getLine[];
    X = ToExpression @ getLine[];
    Do[
        If[500 a + 100 b + 50 c == X, cnt++]
    , {a, 0, A}, {b, 0, B}, {c, 0, C}];
    Print @ cnt;
]
main[]

この問題ではループのために Do を用います。
Do[expr, {i,imin,imax,di}] で、i の値を imin から imax まで di ずつ増やして、順番に expr を評価します。python の range(imin, imax + 1, di) っぽい感じです。また、
- Do[expr, {i,imin,imax}] → Do[expr, {i,imin,imax,1}]
- Do[expr, {i,imax}] → Do[expr, {i,1,imax,1}]
です。多重ループは Do[expr, {i,imin,imax,di}, {j,jmin,jmax,dj}, ...] とするだけです。評価順は i → j です(for (i in ...) { for (j in ...) { expr } })。

ループさえ分かればやるだけですね。一応、lhs == rhs は Equal[lhs, rhs] のシンタックスシュガーです。

第 5 問: ABC 083 B - Some Sums

5.wls
main[] := Module[{lines, lineno = 1, getLine, N, A, B, n, ans = 0, s},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    {N, A, B} = ToExpression /@ StringSplit @ getLine[];
    Do[
        s = Total @ IntegerDigits @ n;
        If[A <= s <= B, ans += n]
    , {n, N}];
    Print @ ans;
]
main[]

「10 進法での各桁の和」を求める必要があります。IntegerDigits を用いると、10進法での桁数字の List が得られます (IntegerDigits[123] で {1, 2, 3})7。
あとは Total で各桁の和が得られます。

比較は lhs <= rhs で LessEqual[lhs, rhs] になります。A <= s <= B のようにつなげると、A<=s && s<=B とほぼ同じ意味になります(e1 && e2 は And[e1, e2])。ちなみに、e1 <= e2 == e3 > e4 != e5 のようなこともできます。

第 6 問: ABC 088 B - Card Game for Two

6.wls
main[] := Module[{lines, lineno = 1, getLine, N, a, s = 0},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    N = ToExpression @ getLine[];
    a = ToExpression /@ StringSplit @ getLine[];
    a = Sort[a, Greater];
    Print[Total @ Take[a, {1, -1, 2}] - Total @ Take[a, {2, -1, 2}]];
]
main[]

Sort を用いてソートします。Sort[list, p] で、比較関数 p を用いて list をソートします。Greater[x, y] は x > y と等価なので、Sort[list, #1 > #2 &] と同じ意味ですね。

Take は基本的には Take[list, n] で list の先頭 n 個を取ってくる関数ですが、Take[list, {min,max,d}] で「min 番目から max 番目まで d 刻みでとってくる」ことができます。-1 は末尾を表します。つまり、Total @ Take[a, {1, -1, 2}] で a の奇数番目の総和が求まります。

第 7 問: ABC 085 B - Kagami Mochi

7.wls
main[] := Module[{lines, lineno = 1, getLine, N, d, i},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    N = ToExpression @ getLine[];
    d = ConstantArray[0, N];
    Do[d[[i]] = ToExpression @ getLine[], {i, N}];
    Print @ Length @ Union @ d;
]
main[]

ConstantArray[c, n] は c を n 個複製した List を生成します。

Union は重複を取り除く関数です8。ついでにソートもしてくれます。Length は要素の数を与えます9。

第 8 問: ABC 085 C - Otoshidama

Wolfram 言語らしい解法があります。

8.wls
main[] := Module[{lines, lineno = 1, getLine, N, Y, a, b, c, ans},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    {N, Y} = ToExpression /@ StringSplit @ getLine[];
    ans = FindInstance[a + b + c == N && 10000 a + 5000 b + 1000 c == Y && a >= 0 && b >= 0 && c >= 0, {a, b, c}, Integers];
    Print @ If[Length[ans] == 0, "-1 -1 -1", StringRiffle[{a, b, c} /. ans[[1]]]]
]
main[]

FindInstance を使うとスマートに解けます。FindInstance[expr, vars, dom] は、expr が成り立つよな変数 vars の割り当てを探索します(dom は探索範囲で、今回は整数なので Integers)。もし解がなければ、空の {} が返され、解が見つかった場合は、vars を {x1, x2, ...} として、{{x1->v1, x2->v2, ...}} が返されます。これは、expr に現れる x1 を v1 に、x2 を v2 に、…、置き換えたときに、expr が True になることを表します。
例えば、FindInstance[x^2 - 3 y^2 == 1 && 10 < x < 100, {x, y}, Integers] は {{x -> 26, y -> 15}} を返しますが、実際に x に 26、y に 15 を代入すると成り立っていることがわかります10。ちなみに、lhs -> rhs は Rule[lhs, rhs] のシンタックスシュガーです。

/. は何でしょうか。expr /. rules は ReplaceAll[expr, rules] のシンタックスシュガーです。rules はルール lhs->rhs、またはルールの List で、expr に現れる lhs をすべて rhs に置き換えます。
例えば、ans = {{a -> 0, b -> 9, c -> 0}} となっていたら、{a, b, c} /. ans[[1]] は

{a, b, c} /. {a -> 0, b -> 9, c -> 0}
{0, 9, 0}

と評価されます。

第 9 問: ABC 049 C - Daydream

正規表現が使えるので、やるだけです。

9.wls
main[] := Module[{lines, lineno = 1, getLine, S},
    getLine[] := lines[[lineno++]];
    lines = StringSplit[$ScriptInputString, "\[NewLine]"];
    S = getLine[];
    Print @ If[StringMatchQ[S, RegularExpression @ "^(dream|dreamer|erase|eraser)*$"], "YES", "NO"];
]
main[]

RegularExpression で正規表現オブジェクトを作ります。

StringMatchQ でマッチするか判定します。

第 10 問: ABC 086 C - Traveling

10.wls
main[] := Module[{blocks, N, d},
    blocks = ImportString[$ScriptInputString, "Table"];
    N = blocks[[1, 1]];
    blocks[[1]] = {0, 0, 0};
    blocks = Transpose @ blocks;
    d = Differences @ blocks[[1]] - Abs /@ Differences @ blocks[[2]] - Abs /@ Differences @ blocks[[3]];
    Print @ If[Min[d] >= 0 && AllTrue[d, EvenQ], "Yes", "No"];
]
main[]

wolframscript コマンドはおそらく Wolfram Kernel を起動していて、それだけで 1 秒ちょっと使います。また、入力の読み込みにもかなり時間がかかってしまい、ロジック部分は 0.1 秒もかかっていないのに、N=10^5 だと、全体では 2.3 秒程度かかり、TLE です。誰か助けてください。

とりあえず解説をします。ImportString[str, "Table"] は str を表として解釈します。サンプル入力の

2
3 1 2
6 1 1

であれば、{{2}, {3, 1, 2}, {6, 1, 1}} になります。ちなみに、Wolfram 言語では行列を List の List で表現します。

list[[i, j]] は list[[i]][[j]] と等価です(list[[i, j, k, ...]] も同様)。

Transpose は行列を転置します。
すると、blocks[[1]] = {0, 0, 0}; blocks = Transpose @ Drop[blocks, 1] で blocks は {{0, 3, 6}, {0, 1, 1}, {0, 2, 1}} になります。blocks[[1]] が t に、blocks[[1]] が x に、blocks[[1]] が y に対応します。
あとは dt - dx - dy が非負の偶数であることを確かめればよいです。
ループを書くと遅いので、Map や Differences, AllTrue を用いて工夫します。

Differences は List の差分を計算します。Differences[{a, b, c, ...}] なら {b-a, c-b, ...} です。Abs は絶対値を計算する関数です。

d = Differences @ blocks[[1]]
    - Abs /@ Differences @ blocks[[2]]
    - Abs /@ Differences @ blocks[[3]]

これで、d = dt - dx - dy が計算できました。AllTrue[list, test] は list のどの要素に test 関数を適用しても True になるかを調べます。

まとめ

正直、Wolfram Kernel の起動が遅すぎるのと、入力が遅いせいで競プロではかなり不利です (幾何には強そう?)。

しかし、他にはない強力な関数 (FindInstance、Resolve、GroebnerBasisなど) が数多くあり、記号計算、任意精度実数が容易に扱え、項書き換えシステムという独特の文法を持っていて、面白い使い方が無限にできそうな可能性を感じます。また、今回は紹介しませんでしたが、C や C++、C# などとのインターフェースもあります。が、かなりマイナーな分野だった思います。無償化によってこのあたりも今後活発になるのではないでしょうか。

この記事がきっかけで少しでも Wolframer が増えてくれると嬉しいです。

  1. 実は、ラズパイからなら 無料で使えます！ ↩

  2. ライセンスの制限はもちろんあります。開発のデモやテスト段階なら商用でも使えるらしいですが、デプロイするのは無料ではできないそうです。 ↩

  3. 途中で Return や Throw、Goto を用いるとこの限りではありません。 ↩

  4. あとコメント (*あああ*) (Wolfram言語の文法) ↩

  5. f を省略することもできて、その場合 If[cond, t, Null] と等価です。 また、cond がどちらにも評価されない場合は未評価のままの If[cond, t, f] が返されます。If[cond, t, f, u] を用いると、True でも False でもない場合は u を返します。 ↩

  6. 複数の引数をもつ関数なら #1, #2, ... で参照できます。Function[{x1, x2, ...}, body] では引数の名前は x1, x2, ... になります。 ↩

  7. IntegerDigits[n, b] なら b 進法での n の桁数字が得られます。 ↩

  8. なぜ Union という名前がついているかというと、Union[list1, list2, ...] のようにするとlist1, list2, ... の和集合を計算してくれるからです。 ↩

  9. List に限りません。Length[f[a, b, c]] は 3 になります。arity のようなものです。すべては式なのです。 ↩

  10. a^b は Pow[a,b] のシンタックスシュガーです。a の b 乗を表します。 ↩

5
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
5
3