ardggy's blog

Esc - Meta - Alt - Ctrl - Shift

doctest 駆動 AtCoder

AtCoder を始めた。ためしに Haskell や Nim で解いたりしていたが、結局は Python3 に落ち着いた。

解くのはいいが、サンプルのテストが面倒でしょうがない。TDD でできるといいが、フレームワークほどの仕組みは必要なさそう。その点で doctest はよかった。プログラムに埋め込めて、テスト実行も $ python -mdoctest XXX.py で済む。エイリアスalias t='python -mdoctest' としておけば t XXX.py だけでテストが走る。

テンプレートはいろいろ調整して以下のようになった。 as_input の仕組みは他の Pythonista の回答を参考にさせていただいた。main 関数の引数がいびつだが、Fail 時にどのケースで失敗したかがわかるのでこうしている。このまま提出しても問題ないようになっている。

#!/usr/bin/env python

def main(_=0):   # 引数はとくに使わない。`as_input` を受け入れるためだけ
    N = int(input())
    print(N * (N + 1) // 2)

# string でテストケースを与えると stdin のように振る舞ってくれる
def as_input(s: str) -> None:
    import io
    global input
    f = io.StringIO(s)
    input = lambda: f.readline().rstrip()
    return None

sample1 = """5
"""
sample2 = """100
"""
def test():
    """
    >>> main(as_input(sample1))
    15
    >>> main(as_input(sample2))
    5050
    """
    pass

if __name__ == '__main__':
    main()

結果は (doctest -v) でこんなかんじにでてくれる。

$ alias t='python -mdoctest'
$ t -v test.py
Trying:
    main(as_input(sample1))
Expecting:
    15
ok
Trying:
    main(as_input(sample2))
Expecting:
    5050
ok
3 items had no tests:
    test
    test.as_input
    test.main
1 items passed all tests:
   2 tests in test.test
2 tests in 6 items.
2 passed and 0 failed.
Test passed.

解が浮動小数点の場合、とりあえずサンプルは実行するが、どうせ間違っているので目視でなんとなくあっていれば出している(適当)

複数のイテレータを合流させる(承前)

前は2つのイテレータを合流させるパターンを書いたが、3つ以上のイテレータも混ぜられないだろうか。 confluence(confluence($it1, $it2, $compare), $it3, $compare) とやれば複数混ぜられそうな気がする。

つまり畳み込みの計算で、array_reduce をつかえばいい。

<?php
declare(strict_types=1);

// x=0: 0, 3, 6, 9, ...
// x=1: 1, 4, 7, 10, ...
// x=2: 2, 5, 8, 11, ...
function mod3($x = 0): Iterator {
  for ($y = $x; 1; $y += 3) { yield $y; }
}

$compare = fn($x, $y) => $x <=> $y;
$iter3 = array_reduce([mod3(0), mod3(1), mod3(2)],
                      fn($acc, $it) => confluence($acc, $it, $compare),
                      $initial = new EmptyIterator);

assert([0,1,2,3,4,5,6,7,8,9] === iterator_to_array(new LimitIterator($iter, 0, 10)));

array_reduce は左からの畳込みだけど、右から畳み込んでも動くんじゃないかと思う。

イテレータを合流させる

ソート済みの2つのイテレータを混ぜ合わせて直列にできないだろうか、という動機

<?php
declare(strict_types=1);

$iter = confluence(odds(), evens(), fn($x, $y) => $x <=> $y);
assert([0,1,2,3,4,5,6,7,8,9] === iterator_to_array(new LimitIterator($iter, 0, 10), $use_keys = false));

function confluence(Iterator $it1, Iterator $it2, Callable $compare): Iterator {
  while ($it1->valid() && $it2->valid()) {
    $v1 = $v1 ?? $it1->current();
    $v2 = $v2 ?? $it2->current();

    // どちらのイテレータからもらうかの選択
    // 改善の余地あり
    if ($compare($v1, $v2) < 0) {
      yield $v1;
      $it1->next();
      $v1 = $it1->current();
    } else {
      yield $v2;
      $it2->next();
      $v2 = $it2->current();
  }
  for ( ; $it1->valid(); $it1->next()) {
    yield $it1->current();
  }
  for ( ; $it2->valid(); $it2->next()) {
    yield $it2->current();
  }
}

function odds(): Iterator {
  for ($x = 1; 1; $x += 2) {
    yield $x;
  }
}

function evens(): Iterator {
  for ($x = 0; 1; $x += 2) {
    yield $x;
  }
}