Arantium Maestum

プログラミング、囲碁、読書の話題

PythonでForth処理系を書く! その7(繰り返し処理 DO-LOOP)

繰り返し処理を実装する。

PythonでForth処理系を書く! その2(実装する機能) - Arantium Maestumでも書いた通り、この構文の要素は3つあって:

  • DOでスタックに乗っている二つの整数を分岐のために使う

    • スタックから得たはじめの整数がループカウンタの開始値
    • スタックから得た次の整数がループカウンタの終了値
  • LOOPでループカウンタをインクリメントする

    • ループカウンタが終了値未満の場合はDO直後のコマンドまで後方ジャンプ
    • ループカウンタが終了値に達した場合はLOOP以降のコマンドを実行
  • DOとLOOPの間で使われるIは現在のループカウンタの値をスタックに乗せる

こんな使い方をする:

0 11 1 DO I + LOOP . CR

上記のコードは1以上11未満の整数をループ足し合わせて55を出力して改行する。

コード

github.com

IF-ELSE-THENと同じようにcompilerとinterpreterに変更を加える。

compiler.py

コンパイラにDO、LOOP、Iトークンに対応する関数を追加する。

DO:

def run_do(state):
    state['branch'].append(len(state['codes']) + 1)
    state['codes'].append('DO')

codesにはDOだけ載せ、branchスタックにそのあとの位置を載せる。(ループする場合に戻るべきジャンプ先)

LOOP:

def run_loop(state):
    do_pos = state['branch'].pop()
    state['codes'].extend(('LOOP', do_pos))

DOのコンパイルでbranchスタックに載せられた「ジャンプ先」をとり、LOOPコードと一緒にcodesに載せる。

I:

def run_i(state):
    state['codes'].append('I')

Iはそのままcodesに載せるだけ。

このとおり、compilerへの変更は非常に簡単。ループ条件はメインスタック上で管理されるので、処理のロジックはinterpreter側で実装されている。

interpreter.py

まずインタプリタの内部状態に新しくbranchスタックを追加する:

def init_state():
    return {
            'codes': [],
            'pos': 0,
            'stack': [],
            'branch': [], #新しく追加
            'end': False,
            }

IF-ELSE-THENの時にコンパイラには追加したが、今回はインタプリタにも必要。

ループカウンタの現在の値と終了値を保持するために使う。この二つをメインスタックにいれておくと、ループ中にスタックに値を追加するなどの処理と競合してしまう恐れがある。

DOの処理:

def run_do(state):
    a, b = state['stack'].pop(), state['stack'].pop()
    state['branch'].extend((b, a))
    state['pos'] += 1

スタックから二つの数字を取り、片方を開始時のループカウンタ、もう片方をループの終了値としてbranchスタックに入れる。

LOOPの処理:

def run_loop(state):
    a, b = state['branch'].pop(), state['branch'].pop()
    a += 1
    if a < b:
        state['branch'].extend((b, a))
        state['pos'] = state['codes'][state['pos']+1]
    else:
        state['pos'] += 2

branchスタックからループカウンタと終了値をとり、ループカウンタをインクリメントしてから比較。

  • ループカウンタが終了値以上なら終了
  • それ以外なら
    • ループカウンタ・終了値をbranchスタックに戻す
    • コンパイルされているループ元の位置にジャンプ

Iの処理:

def run_i(state):
    state['stack'].append(state['branch'][-1])
    state['pos'] += 1

branchスタックにあるループカウンタをメインスタックにも載せる(branchスタック側の値は載せたままで消費しない)。これでカウンタの値を利用した計算などができる。

使ってみる

まずはDOとLOOPだけで試してみる:

Starting Forth REPL...
> 5 0 DO 1 . LOOP CR
11111
> 0 5 DO 1 . LOOP CR
1

DO-LOOPは必ず一回は処理が走るので、0 5 DO 1 . LOOP CRのようにカウンタの開始値が終了値よりも大きい場合でも1 .の部分が一回実行されているのがわかる。

次にIも使ってみる:

> 5 0 DO I . CR LOOP
0
1
2
3
4
> 0 11 1 DO I + LOOP . CR
55
> 0 10 0 DO I 1 + + LOOP . CR
55

ループカウンタの値がちゃんと計算に使われている。

ちなみに0 10 0 DO I 1 + + LOOP . CRはループカウンタが0からはじまり10未満までループするなか、カウンタの値Iをインクリメントしてからこれまでの合計に足し合わせている。

次回は変数定義。

PythonでForth処理系を書く! その6(条件分岐 IF-ELSE-THEN)

条件分岐を実装する。

PythonでForth処理系を書く! その2(実装する機能) - Arantium Maestumでも言ったとおり、Forthの条件分岐構文は他の言語を知る人には多分だいぶ奇妙に見えるはず。

上記の記事からの再掲になるが、このような構文で:

1 IF 2 ELSE 3 THEN . CR

このような処理になる:

IFコマンドの時点でスタックの上が0以外ならELSEまで実行してからTHENにジャンプ、0ならELSEまでジャンプしてから実行を続ける

THENが「IFが満たされた場合に実行される処理」ではなく、「IF-ELSE分岐が終わって実行パスが合流するポイント」を表しているのがForthの特徴。THENではなくENDIFのほうがいいのではないかという意見も読んだことがあり、一理あると思う。

実装としてはこう:

github.com

ついでにこう(バグの修正):

github.com

変更している箇所はcompilerとinterpreter。ここでようやく中間言語コンパイルすることの利点が活きてくる。

compiler.py

まずはコンパイラの内部状態にあたらしく分岐に関連する情報を維持するスタックを用意:

def init_state():
    return {
            'tokens': [],
            'pos': 0,
            'codes': [],
            'branch': [],  #これを追加
            'end': False
            }

そのbranchスタックを使ってIF、ELSE、THENのトークンをコンパイルする各関数を定義する。

まずIF:

def run_if(state):
    state['branch'].append(len(state['codes']) + 1)
    state['codes'].extend(('IF', None))

出力するcodesにIFとNoneを載せる。

コンパイルが進んで対応するELSE(あるいはELSEがない場合はTHEN)に遭遇したときに、NoneをそのELSE/THENのコードの位置の値に書き換える。あとで書き換えられるようbranchスタックにNoneの位置(len(state['codes']) + 1)を載せておく。

次にELSE:

def run_else(state):
    if_pos = state['branch'].pop()
    state['codes'][if_pos] = len(state['codes']) + 2

    state['branch'].append(len(state['codes']) + 1)
    state['codes'].extend(('ELSE', None))

まずはbranchスタックから対応するIFのNoneの位置を取り出す。そのNoneをELSE以降のコードの位置(len(state['codes']) + 2)に書き換える。

ELSEで追加されるコードはIFと似たようなもの。ELSEとNoneで、NoneはELSEと対応するTHENのコードにジャンプするための位置情報。

最後にTHEN:

def run_then(state):
    else_pos = state['branch'].pop()
    state['codes'][else_pos] = len(state['codes'])

ELSEの前半部分とほぼ同一。branchスタックに載っているのが対応するELSEのNoneの位置(あるいはELSE部分がなくIF-THENだった場合はIFのNoneの位置)。

THENの場合はそこに到達した時点で分岐は終了しており、ジャンプする必要がないのでcodesにはなにも追加されない。

最後に定義した関数をprimitivesとして追加:

primitives = {
        '+': run_plus,
        '.': run_print,
        'CR': run_cr,
        'IF': run_if,
        'ELSE': run_else,
        'THEN': run_then,
        }

IF/ELSE/THENトークンに遭遇したらrun_if/run_else/run_then関数が呼び出されるようになる。

実はここでバグをいれて放置してしまった。「THENがコード追加しない」ということから何故か「primitivesに追加しなくてもオッケー」と思い込んでいて、最初のbranch mergeでは'THEN': run_thenの部分が抜けていた。不覚・・・

このコンパイラ

1 IF 2 ELSE 3 THEN . CR

このコードをコンパイルすると

['PUSH', 1, 'IF', 8, 'PUSH', 2, 'ELSE', 10, 'PUSH', 3, '.', 'CR', 'END']

こんな感じの中間言語に変換される。あるいはコード位置も追加でつけるなら:

 0 PUSH
 1 1
 2 IF
 3 8
 4 PUSH
 5 2
 6 ELSE
 7 10
 8 PUSH
 9 3
10 .
11 CR
12 END

PUSHの後の数字はスタックに載せるための数字、IFやELSEの後の数字はジャンプ先の位置だ。

ちなみにTHENトークンの情報はELSEのジャンプ先の数字になるので、THENというコード自体はコンパイル結果には現れないことがわかる。

次にこれをインタプリタが実行できるようにする。

interpreter.py

IFとELSEのコードを処理できるようにinterpreterに関数を追加する。

まずはIF:

def run_if(state):
    a = state['stack'].pop()
    if a == 0:
        state['pos'] = state['codes'][state['pos']+1]
    else:
        state['pos'] += 2

スタックの一番上をとり、その数が0であればジャンプ、そうでなければそのまま次のコードを処理していく。

IFのすぐあとはジャンプ先の位置情報なので、もしジャンプしない場合はそのコードも飛ばす必要があるのでstate['pos'] += 2になる。

ELSEの処理:

def run_else(state):
    state['pos'] = state['codes'][state['pos']+1]

IFの条件が0だった場合はELSEの先にジャンプするので、ELSEコードに到達したということはすでにIFからELSEの間のコードが走ったということ。なので、ELSEからTHENまでのトークンがコンパイルされた「ELSEブロック」を飛ばしてその先にジャンプする。

使ってみる

前回実装したREPLでいろいろといじっていみる:

Starting Forth REPL...
> 1 IF 2 ELSE 3 THEN . CR
2
> 0 IF 2 ELSE 3 THEN . CR
3
> 1 1 IF 2 THEN . CR
2
> 1 0 IF 2 THEN . CR
1
> 1 -1 + IF 0 ELSE 2 THEN . CR
2

ELSEが存在する場合、存在しない場合どちらでもうまく作動している。

IF-ELSE-THENのネスト:

> 0 1 IF IF 2 ELSE 3 THEN ELSE 4 THEN . CR
3
> 1 1 IF IF 2 ELSE 3 THEN ELSE 4 THEN . CR
2
> 0 0 IF IF 2 ELSE 3 THEN ELSE 4 THEN . CR
4

ややこしいが正しい挙動になっている。

IF-ELSE-THENの順序が間違っていたりするとエラー:

> 1 IF 5 THEN 3 ELSE 2 . CR
Traceback (most recent call last):
  File "main.py", line 36, in <module>
    run_repl()
....
    if_pos = state['branch'].pop()
IndexError: pop from empty list

異常系でどんな挙動になるかはまったく保証できない。今回はコンパイルエラー(ELSEが書き換えるべきIFのジャンプ先がbranchスタックに乗っていなかった)。今のところエラーハンドリングは改善の予定はない。めんどくさそう。

次回はDO-LOOP構文を追加して「繰り返し」処理を可能にする。

PythonでForth処理系を書く! その5(Read-Eval-Print-Loop)

REPL機能を実装する。

main.pyに引数がなければREPLを実行するようにする(引数としてファイル名が渡されていればそのファイルを実行)。REPLは一行ずつユーザからForthコードを受け取り、今まで受け取ったコードの環境を引き継ぎつつ実行する。

github.com

変更はcompiler、interpreter、そしてmainの三箇所。

compiler.py

compile関数が複数回呼ばれることを前提に以下の変更を加える:

  • 中間言語に変換するtokensを初期化せずに追加するようにする
  • 中間言語のcodesから'end'コードを取り除く
  • 'end'フラグをFalseにセットする
def compile(state, tokens):
    state['tokens'].extend(tokens)
    if state['codes']:
        state['codes'].pop()
    state['end'] = False

    while not state['end']:
        run(state)
    return state['codes']

どのトークンまで変換したかを表す'pos'値は以前の最後のトークンの一つ先を指しているので、ちょうど新たに追加したトークンの先頭を指すことになる。

なのでこれでcompile関数が走れば追加分のトークンだけ中間言語に変換してくれる。

interpreter.py

変更はinterpret関数のstate['end']をFalseにセットするだけ。

def interpret(state, codes):
    state['codes'] = codes
    state['end'] = False #ここだけ変更

    while not state['end']:
        run(state)

main.py

  • 今までのmain部分をrun_fileという関数に
  • run_repl関数を別途定義
  • main部分ではスクリプトが呼ばれたときにargvに引数が渡されたかどうかでrun_fileかrun_replにディスパッチ

今回追加したrun_repl関数はこんな感じ:

def run_repl():
    compiler_state = compiler.init_state()
    interpreter_state = interpreter.init_state()

    print("Starting Forth REPL...")

    while True:
        s = input('> ')
        if s == 'QUIT': break

        tokens = tokenizer.tokenize(s)
        codes = compiler.compile(compiler_state, tokens)
        interpreter.interpret(interpreter_state, codes)

まずは内部状態を初期化して、ループ内で新しくユーザから受け取った文字列をトークン化し、内部状態と一緒に引数としてcompile/interpretに渡して処理している。

実行

これで単にpython main.pyと呼び出せば:

Starting Forth REPL...
> 3 4 +
> 5 6 +
> + . CR
18

というような感じで対話的にForthを書ける。行の間で環境(とくにスタックの状態)が引き継がれているのがわかる。

これで機能を追加したときにチェック・デバッグしやすくなった。

次回はI分岐構文のIF-THEN-ELSEを追加する。

PythonでForth処理系を書く! その4(内部状態をPython Dictionaryに保存)

前回の終わりで書いたように、Forth処理系の内部状態、とくにコンパイラ部分とインタプリタ部分のものをデータ構造として抜き出して、保存・出力できるようにしたい。

その変更がこれ:

github.com

主要な変更は二点:

  1. 内部状態をPythonのDictionaryで表し、compileやinterpretにそのState dictionaryを渡している
  2. コンパイラインタプリタのどちらにも存在したif-elif-elseの分岐処理を、個別にState dictionaryを受け取ってひとつの処理を行う複数の小さい関数に分割する

state dictionaryの導入

とくにREPLを実装するのに役にたつと考えての変更。

一つのプログラムを最初から最後まで一気に読み込んで処理を行う場合は、こういう状態が関数内で完結していて問題ないのだが、ユーザから一行ずつプログラムを与えられて、その都度「トークン化・コンパイル・実行」するとなると、過去の状態も含めて保持する必要がある。

その方法としては

  • オブジェクト
  • Python generator
  • グローバル変数としてのState
  • 引数・戻り値としてのイミュータブルState
  • 引数としてのミュータブルState

あたりが考えられると思う。

一番素直なのはオブジェクト化してしまって状態をそのままメンバ変数にいれてしまうことだろう。あるいはPythonだとyieldキーワードを使ったgeneratorで関数に状態を手軽に持たせることもできる。ただ、デバッグなどの関係もあって、できれば状態を簡単なデータ構造に入れて外部からアクセスできるようにしたかった。(もちろんオブジェクトでもメソッドを定義するなどして実現できるが少し手間だ・・・)

グローバル変数は宗教上の理由でアウト。イミュータブルなStateを引数・戻り値にするのは好みにかなっているのだが、Pythonのデータ構造はそもそもmutationを前提としていてイミュータブルに使うにはコピーコストが大きい。

というわけでcompileやinterpretに引数としてStateを与えて、それに対して破壊的変更を加えていく形のコードになった。

if-elif-elseの分割

Stateをデータ構造につめこんだので、一つの関数内で処理が完結する必要はなくなった。

  • 分岐で行っていた各処理を個別の関数にわけて
  • 現在位置にあるトークン・コードによってディスパッチ
  • Stateも引数として与えて変更がかかる

という流れにする。このような実装にすることで今後機能を追加していくときに「一つのでかい条件分岐にelifを追加してどんどん長くしていく」という地獄を避けられるはず。

次回は今回の変更を踏まえてインタプリタにREPL機能をつける。

PythonでForth処理系を書く! その3(最低限の機能)

PythonでForth処理系を書く! その2(実装する機能) - Arantium Maestum

でも書いた通り、まずは以下の機能を最低限として実装する:

  • 整数をスタックに乗せる
  • スタックから二つの整数をとり、それらを足し合わせた整数をスタックに乗せる
  • スタックから整数をひとつとり、その数を標準出力に表示する
  • 標準出力に改行を表示する

実行できるコードはこんな感じだ:

1 2 + . CR

簡単な実装

まずはできるだけ簡単に実装してみる:

s = '1 2 + . CR'

stack = []
for token in s.split():
    try:
        stack.append(int(token))
        continue
    except ValueError:
        pass
    if token == '+':
        a, b = stack.pop(), stack.pop()
        stack.append(a+b)
    elif token == '.':
        print(stack.pop(), end='')
    elif token == 'CR':
        print()

ハードコードされたforthコードをsplitでトークンに分け、forループで一つ一つのトークンごとにtry-except/if-elseでトークンタイプをチェックして実行していく。

今回実装する機能だけだったらこれで充分かもしれない。ただ、密結合な上にフローをかなり単純視している(具体的にはトークンとインタプリタの動作に1対1の対応があることを前提としてしまっている)ので、今後の機能追加が容易ではない。

拡張しやすくなるように、どのような異なる機能・段階があるのかを明確にする形で分割してみる。

分割した実装

github.com

まずはインタプリタの全体を統括するmain.py:

import tokenizer
import compiler
import interpreter

if __name__ == '__main__':
    with open('samples/sample1.forth') as f:
        s = ' '.join(f)
    tokens = tokenizer.tokenize(s)
    code_list = compiler.compile(tokens)
    interpreter.interpret(code_list)

分割した個別のトークナイザ・コンパイラインタプリタをインポートして使っている。一つのポイントとしては、全体としてはインタプリタなのだが、トークンを直接インタプリタに渡すのではなく、インタプリタが実行しやすいような中間言語コンパイルするステップが入っていること。このコンパイルステージが今後の条件分岐や変数定義などでは非常に便利になる。

tokenizer.pyの実装:

def tokenize(s):
    return s.split()

非常に簡単。スペースで分割するだけ。トークンも特にクラス化せず、ただの文字として扱う。

compiler.py:

def compile(tokens):
    code_list = []
    for token in tokens:
        code_list.extend(compile_token(token))
    code_list.append('END')
    return code_list

def compile_token(token):
    if is_int(token):
        return ('PUSH', int(token))
    d = {'+':'PLUS', '.':'PRINT', 'CR':'CR'}
    return (d[token], )

def is_int(s):
    try:
        int(s)
    except ValueError:
        return False
    return True

トークンが数字だったらスタックに乗せるようにPUSH Xという二つのコードにコンパイル。それ以外(+, ., CR)はそれぞれ一つのコードにコンパイルする。

トークンが数字かどうかの判定もコンパイラ部分で行っているが、トーケナイザ部分で判定してその情報をトークンに持たせる手もある。あまり関係ないが、個人的にPythonの「文字列がInt化できるか判定」イディオムはあまり好きではない・・・

コンパイラの吐く中間言語を受け取り、実行していくinterpreter.py:

def interpret(code_list):
    counter = 0
    stack = []
    while True:
        if code_list[counter] == 'END':
            break
        elif code_list[counter] == 'PUSH':
            stack.append(code_list[counter+1])
            counter += 2
        elif code_list[counter] == 'PLUS':
            a, b = stack.pop(), stack.pop()
            stack.append(a+b)
            counter += 1
        elif code_list[counter] == 'PRINT':
            print(stack.pop(), end='')
            counter += 1
        elif code_list[counter] == 'CR':
            print()
            counter += 1

スタック(Pythonのリスト)に数字をappendしたりpopしたりしながら中間言語の指示を実行していく。

これで元々のサンプル:

1 2 + . CR

を実行することはもちろん、今後の機能を追加していくのもかなりやりやすくなった。だが、実はもう一つリファクタするべきポイントがある。コンパイラインタプリタの両方で、状態をデータ構造に保存するようにしたい。

次回はその理由と実装の話。

PythonでForth処理系を書く! その2(実装する機能)

処理系で実装する機能を列挙してみる。

  • 整数をスタックに乗せる
  • スタックから二つの整数をとり、それらを足し合わせた整数をスタックに乗せる
  • スタックから整数をひとつとり、その数を標準出力に表示する
  • 標準出力に改行を表示する
  • REPL機能
  • 分岐構文 IF-ELSE-THEN
  • ループ構文 DO-LOOP
  • 変数定義
  • 関数定義

第1段階 - 最低限の機能

まず手始めに実装するもの:

  • 整数をスタックに乗せる
  • スタックから二つの整数をとり、それらを足し合わせた整数をスタックに乗せる
  • スタックから整数をひとつとり、その数を標準出力に表示する
  • 標準出力に改行を表示する

以上で例えばこんなコードをファイルから読み出して実行できるようになる:

1 2 + . CR

処理としては

  • 1をスタックに乗せる
  • 2をスタックに乗せる
  • 2と1をスタックから取り出し、足した結果の3をスタックに乗せる
  • 3をスタックから取り出し、出力する
  • 改行を出力する

という流れになる。

第2段階 - REPL

上記の機能をインタラクティブなREPLを通して使えるようにする。

具体的には、一行入力されるごとにコードをインタプリタが実行していく。スタックの状態は行をまたいで維持される。

第3段階 - 分岐構文

IF-ELSE-THEN構文を実装する。

この構文は(とくにTHENが)ちょっと独特で、このように使う:

1 IF 2 ELSE 3 THEN . CR

どういう論理かというと

IFコマンドの時点でスタックの上が0以外ならELSEまで実行してからTHENにジャンプ、0ならELSEまでジャンプしてから実行を続ける

というもの。

上記のコードは2を出力してから改行する。 一番先頭の1が0なら「3を出力してから改行」になる。

第4段階 - ループ構文

DO-LOOP構文を実装する。

この構文の要素は3つあって:

  • DOでスタックに乗っている二つの整数を分岐のために使う スタックから得たはじめの整数がループインデックスの始点 スタックから得た次の整数がループインデックスの終点

  • LOOPでループインデックスをインクリメントする ループインデックスが終点未満の場合はDO直後のコマンドまで後方ジャンプ ループインデックスが終点に到達した場合はLOOP以降のコマンドを実行

  • DOとLOOPの間で使われるIは現在のループインデックスの値をスタックに乗せる

こんな使い方をする:

0 11 1 DO I + LOOP . CR

上記のコードは1以上11未満の整数をループ足し合わせて55を出力して改行する。

第5段階 - 変数定義

Forthでの変数はVARIABLEで宣言、!で代入、@で参照となる。

VARIABLE X
VARIABLE Y
5 X !
10 Y !
X @ Y @ + . CR

のように使う。上記のコードを走らせると15が出力される。

まずVARIABLE X、VARIABLE YでXとYを未使用のメモリ領域に結びつける。VARIABLE宣言の後にはXやYというコードはそのメモリ領域のアドレスをスタックに乗せるコマンドになる。

5 X !では5をスタックに乗せ、Xのアドレスをスタックに乗せ、そして!でその二つの整数をとり、最初の整数をアドレスとして次の整数をそのアドレスの値として代入する。

X @でXのアドレスをスタックに乗せ、そのアドレスをとってその領域に入っている値をスタックに乗せる。

第6段階 - 関数定義

Forthでは特定の処理が紐付いたコードのことをワードと呼ぶ。関数もワードだ。

ユーザがワードを自前で定義できるようにするための構文が:と;だ。

:がワード宣言の冒頭、;がワード宣言の終了を示す。:の次のコードは新しいワード名、それ以降で;までのコードはそのワードに紐づけられる処理だ。

: ADD2 2 + ;
1 ADD2 . CR

で3が出力される。

: ADD2 2 + :でADD2というワードがその処理2 +とともに登録される。

1 ADD2 . CRはADD2の部分で登録された処理が呼び出されるので結果としては1 2 + . CRと同じ処理となる。

その後

ここまで実装してとりあえず終わる予定。

この枠組みに四則演算やスタック上の値を操作(コピー・廃棄・上位Xを入れ替えなど)するような機能で肉付けしていくと、結構簡単に非常に表現力の高い言語処理系ができあがるようだ。こっちの実装はForthについてもっと理解が進んだら試してみたい。

AtCoder Beginner Contest 094の問題を解いた

第1問 Cats and Dogs

abc094.contest.atcoder.jp

abc094.contest.atcoder.jp

猫がA匹いることがわかっているのに加えて、追加で猫か犬かがB匹いる場合、猫がX匹いることが可能か。

不可能になるのはX < AかX > A + Bの二ケースなので、A <= X <= A+Bが成り立っているかどうかを調べる:

a, b, x = map(int, input().split())
print('YES' if (a <= x <= a+b) else 'NO')

第2問 Toll Gates

abc094.contest.atcoder.jp

abc094.contest.atcoder.jp

直線上に順番にならんでいる料金所を通るたびにコスト1かかるとして、地点Xからはじめて直線の左右どちらかの端に到達するのにかかる最小コストを求める。

料金所がもとからソートされているので、そのまま二分探索が使える。地点X自体には料金所がないことがわかっているのも処理を簡単にしてくれるポイント。

from bisect import bisect_left
 
n, m, x = map(int, input().split())
axs = [int(x) for x in input().split()]
 
i = bisect_left(axs, x)
print(min(i, len(axs) - i))

bisect_left(axs, x)で地点Xの左にいくつ関所があるかがわかる。

右にある料金所の数は「料金所の総数 - 左側にある料金所の数」なのでmin(i, len(axs) - i)で最小コストが求められる。

ちなみにこれ、最初に書いた解だと

print(min(len(axs[:i]), len(axs[i:])))

になっていた。iが持つ意味を直観的に捉えきれてなかったのだろうか。こういうところで何気なく無駄な処理をしてしまうかどうかがけっこうアルゴリズムでは重要な気がする。

第3問 Many Medians

abc094.contest.atcoder.jp

abc094.contest.atcoder.jp

与えられたリストxsのある要素が含まれていない場合の中央値を、各要素ごとに算出していく。

xsの長さが偶数であるという制約で計算が楽になる。

xsの中央値を計算する場合は、中央にある二つの数字m1とm2 (m1 <= m2)の平均をとることになる。

とりあえずxsから最小値を除いたリストzsの中央値を考えてみると、m2であることがわかる。さらに最小値でなくてもm2未満のどの数字を除いたとしてもm2が中央値になる。

そしてm2以上の場合はm1が中央値だ(最大値からはじめて全く同じロジックを考えることができる)。

なので実装はこんな感じ:

n = int(input())
xs = list(map(int, input().split()))
 
ys = list(sorted(xs))
m1 = ys[n//2 - 1]
m2 = ys[n//2]
 
medians = [m2 if x < m2 else m1 for x in xs]
print('\n'.join(map(str, medians)))

実はPython3.4以降statisticsという標準ライブラリが追加され、median_lowとmedian_highという関数が用意されている。

medianはquickselectアルゴリズムを使ったりすると平均でO(N)になったりすることが知られているので、ちょっと期待して使ってみる:

abc094.contest.atcoder.jp

from statistics import median_low, median_high
 
n = int(input())
xs = [int(x) for x in input().split()]
 
m1, m2 = median_low(xs), median_high(xs)
 
medians = (str(m2 if x < m2 else m1) for x in xs)
print('\n'.join(medians))

問題なくACするのだが、もとのコードが200msだったのがstatisticsを使うと300msほどになる。

CPython statisticsモジュールのソースにあたってみると、medianのなかでリストを複製ソートしていた。

元のコードで一番時間がかかるところをmedian_lowとmedian_highで二回やっているわけだからそれは遅くなるわけだ・・・

今度時間があったら自分でquickselectを実装してみたい。

第4問 Binomial Coefficients

abc094.contest.atcoder.jp

abc094.contest.atcoder.jp

nCrを最大化するnとrを与えられた配列xsから選択するという問題。

まず、任意のrで n1 > n2 なら n1 C r > n2 C rであることから、nはxsの最大値。

任意のnで、nCrを最大化するrはn/2に一番近いものだ。

そしてこの問題だとn > rである必要がある(nCn == nC0 == 1なのだが、この設問だと(n,0)がACの場合、(n, n)はWA)

_ = input()
 
axs = list(map(int, input().split()))
 
i = max(axs)
j = min(axs, key=lambda j: abs((i/2) - j) + (i == j))
print('{} {}'.format(i, j))

Pythonのsort/max/minがオプションでkey関数をとるので、「xに一番近い」というのをminで表せる。ついでにi == jの時のペナルティもつけた。ペナルティが必要になるのはリストに最大要素と0しかない場合のみなので、「最大要素の場合1を足す」というペナルティだけで(n, n)になることを防げる。

感想

この回は基礎数学的な素養が試されている印象が強い。典型的なアルゴリズムやデータ構造らしきものを使うことなく、考察が済めばあとはちょっとソートしたりループしたりするだけで、一番アルゴリズムしていたのはB問題の二分探索だった。

頭の体操としては大変面白いのだが、ABC103のように一旦盲点に入ってしまうとアイデアの手がかりがあまりなく思考が空転してしまいそうな怖さがある。どうしたらこの分野を強化できるか(というか安定的に問題解決に向けて思考を重ねていけるか)、ちょっと考えないといけない。