4年前の、そして半年前の簡単な問題を10分で解いてみた

寝る前に、ちょっとしたプログラミング問題を見つけたので挑戦してみた。

本日の目玉を先に、一言で

『データ列を長さnのグループにまとめる常套句』があることを知った。衝撃的だった。

問を提示しているサイト

先駆者たち

問題

あるカードの山を n 人に均等になるように配った時、それぞれのプレーヤーの手札を答えよ。あまったカードは配らないものとする。

カードの山は "1234567" のような文字列で、プレーヤーの人数は整数で、そのカードを n 人に配ったときの答えは ["14", "25", "36"] のような文字列を含むリストで表すこととする。

言い換えると、次のような引数 numPlayers*1, deck を受け取ると以下のリストを返す関数 deal を作成せよ、ということ。

制限時間 10 分。

#引き数
#引き数
#Returns:返り値

3
"123123123"
Returns: ["111","222","333"]

4
"123123123"
Returns: ["12","23","31","12"]

6
"012345012345012345"
Returns: ["000", "111", "222", "333", "444", "555" ]

4
"111122223333"
Returns: ["123", "123", "123", "123" ]

1
"012345012345012345"
Returns: ["012345012345012345" ]

6
"01234"
Returns: ["", "", "", "", "", "" ]

2
""
Returns: ["", "" ]

10分以内を目指して解いてみた

あまりにも時間が短いので動くことを目指してなりふりかまわず前進。そして 8分でとりあえず動くのはできた。けれどひどいコード。そして改良案を考えていたらタイムアップ。

def deal(numPlayers, deck):
    result = []
    for _ in range(numPlayers):
        result.append('')
    for i, card in enumerate(deck[:len(deck)-len(deck)%numPlayers]):
        result[i % numPlayers] += card
    return result

print deal(3, '123123123')
print deal(4, '123123123')
print deal(6, '012345012345012345')
print deal(4, '111122223333')
print deal(1, '012345012345012345')
print deal(6, '01234')
print deal(2, '')

力不足を痛感。

やり直してみた（スライスを使用）

たしか、 Python のスライスって step という第３の引数があったような。

>>> '123456'[0::2]
'135'
>>> '123456'[1::2]
'246'

うん、これだ。あれ？カードがあまってもうまく動くのかコレ？

>>> '1234567'[0::2]
'1357'
>>> '1234567'[1::2]
'246'

ダメだった。当然か。余計なのは事前に削ろう。そうしてできたのがこのようなコード。

def _deal(numPlayers, deck):
    mod = len(deck) % numPlayers
    if mod:
        deck = deck[:-mod]
    for i in range(numPlayers):
        yield deck[i::numPlayers]

def deal(numPlayers, deck):
    return list(_deal(numPlayers, deck))

翌日追記。牌語備忘録さんのコードと似ていたのに気づいたので。牌語備忘録さんの「末尾を削った文字列を作らずにスライスのend引数を使う」のと自分の「事前に末尾の計算を済ませておく」のとのいいとこどりをめざしてみた。

def _deal(numPlayers, deck):
    deck_len = len(deck)
    end = deck_len - deck_len % numPlayers
    for start in range(numPlayers):
        yield deck[start:end:numPlayers]

def deal(numPlayers, deck):
    return list(_deal(numPlayers, deck))

リストの生成に特化させてみた

def deal(numPlayers, deck):
    deck_len = len(deck)
    end = deck_len - deck_len % numPlayers

    result = []
    for start in range(numPlayers):
        result.append(deck[start:end:numPlayers])
    return result

zipzip してみた

一番最初に思いついたアイディアは人数分、たとえば 3 人なら 3 枚ずつに分割できてしまえば zip することでできる、だった。

>>> zip('123', '456', '789')
[('1', '4', '7'), ('2', '5', '8'), ('3', '6', '9')]

だから足りないのは n 枚ずつに分けてくれるジェネレータ。作るのに手間取るといやだから逃げたけれど。制限時間を気にしなくていいのなら、と作ってみた。でもなんか泥臭い。

def separate(iterable, n):
    temp = []
    count = 0
    for item in iterable:
        count += 1
        temp.append(item)
        if not count % n:
            yield tuple(temp)
            temp[:] = []

こんなときは itertools 。あまり理解し切れていないモジュールなだけにオレには思いもよらぬすごいコードが見つかるはず。

そして、ドキュメントに『データ列を長さnのグループにまとめる常套句』なるものを発見。最初に思いついた人やばい！

# s はイテレート可能オブジェクト、 n は整数
itertools.izip(*[iter(s)]*n)

>>> from itertools import izip
>>> izip(*[iter('123456789')]*3)
<itertools.izip object at 0x00AF3CD8>
>>> list(_)
[('1', '2', '3'), ('4', '5', '6'), ('7', '8', '9')]

あまりにも魔法すぎてよくわからないので具体的な値を入れつつ、分解してみた。

>>> iter('123456789') # 一字ずつ返すイテレータ
<iterator object at 0x00AF0210>
>>> [iter('123456789')]*3 # を 3 つ持ったリスト
[<iterator object at 0x00AF01F0>, <iterator object at 0x00AF01F0>, <iterator obj
ect at 0x00AF01F0>]

ここでやっと理解。なんと、リストを乗算すると中のオブジェクトが使いまわされてしまうという普段避けるべき現象を逆手に取っている！なるほどリストの乗算は itertools.repeat 代わりなのか。リストの 3 要素ともに同じイテレータを参照している。

これを zip すると n 個ずつ得られるわけか。

>>> zip(*[iter('123456789')]*3) # 3 つのイテレータ（全部同じ）から一つずつ得る
[('1', '2', '3'), ('4', '5', '6'), ('7', '8', '9')]

これをもう一度 zip にかければタプルで答えが得られるようになる。それを ''.join でくっつけて、と。できた。

from itertools import izip, repeat

def _deal(numPlayers, deck):
    for s in izip(*izip(*[iter(deck)]*numPlayers)):
        yield ''.join(s)

def deal(numPlayers, deck):
    return list(_deal(numPlayers, deck))

しかし問題発生。

>>> deal(6, '01234') # ['', '', '', '', '', ''] が欲しい
[]

原因はコレ。

>>> zip(*[iter('01234')]*6)
[]
>>> zip(*[iter('')]*2)
[]

配るべきカードがないときにも空文字列を返して欲しいのに空のリストになってしまう。配るべきものがないのは山にあるカードの総枚数が人数より少ない時だからそこで分岐することに。ないときは '' を人数分返せばいいから、 [''] *numPlayers か itertools.repeat だ。

そして完成。

from itertools import izip, repeat

def _deal(numPlayers, deck):
    if len(deck) < numPlayers:
        return repeat('', numPlayers)
    return (''.join(s) for s in izip(*izip(*[iter(deck)]*numPlayers)))

def deal(numPlayers, deck):
    return list(_deal(numPlayers, deck))

翌日追記。速度を計ってみたところスライスで分割するコードより遅い。でも timeit(repeat=10000) した時、数百ms 違う程度。一回実行あたり数十マイクロ秒くらい。カードゲームの AI を作るのでもない限り気にする必要はなさそうな速度差だと思う。

*1:個人的には numPlayers より num_players のほうが好み。 PEP 8 もこうなっている。でも、 PEP 8 には、すでに存在するコードとの兼ね合いも考慮して読みやすいほうを「自分」で判断せよ、といった類の文もある。しかも最初に。だから今日は numPlayers 。

銀月の符号

Python 使い見習いの日記・雑記