イテレータブルのある要素より後ろの要素を取り出す (itertools.dropwhile)

Python「インデクサとは違うのだよ、インデクサとは！」 - つまみ食うを読んで。これは itertools モジュールの楽しさを伝える好機、と見たので。

itertools.dropwhile

「ある要素より後ろの要素を取り出す」は「ある要素以前を捨てる」と言い換えることができます。そして、そんな名前を持つ関数が itertools モジュールには存在します。 itertools.dropwhile です。これは「ある条件が成立している間は要素を捨て、初めて不成立となった以降の要素を得る」ジェネレータを作成します。

この dropwhile を次のように用いることで「ある要素以降の要素」が得られます。

>>> from itertools import dropwhile
>>> L = [ "Marc Bolan", "David Bowie", "Mick Ronson",
...       "Ian Hunter", "Morgan Fisher",
...       "Brian Ferry", "Brian Eno", "Phil Manzanera", "Andy Mackay" ]
>>> it = dropwhile(lambda item: item != 'Mick Ronson', L)
>>> it
<itertools.dropwhile object at 0x00AFB8F0>
>>> list(it)
['Mick Ronson', 'Ian Hunter', 'Morgan Fisher', 'Brian Ferry', 'Brian Eno', 'Phil
 Manzanera', 'Andy Mackay']

dropwhile オブジェクトを作った時点では値の評価は行われないため、処理の後回し、たらい回しが可能というメリットがあります。遅延評価なのです。

次に、出力の形式を選ぶことができます。たとえば、今回のように list で受けるとリストにすることができます。 for 文でまわしたり、 tuple で受けたり u', '.join メソッドに与えるなど、さまざまな出力方法を選ぶことができます。

しかし、今回欲しいのは「ある要素より後ろの要素」なので、先頭が余計です。リストを作り上げた後に [1:] スライスで先頭を切っても目的は達成できるのですが、ここではもう一つ道具を追加することで解決してみることにします。 itertools.islice です。

2009/11/4 追記。 dropwhile, islice のあわせ技ではなく takewhile で不要部分を捨てる方法もあります。

itertools.islice

itertools.lslice は「イテレータから要素を選んで取り出す」ジェネレータを作成します。使い方はスライスに酷似しています。

>>> from itertools import islice
>>> islice(u'abcdefghijklmn', 1, 6, 2)
<itertools.islice object at 0x00AF96C0>
>>> u''.join(_)
u'bdf'

この例では u'abcdefghijklmn'[1:6:2] を回りくどい書き方にしただけではありますが。

islice も dropwhile 同様、作った時点では評価は行われません。

今回は最初の一つを捨てるために使うので、 [1:] スライスを模倣すれば解決です。これを実現する引数は 1, None です。

>>> islice(u'abcdefghijklmn', 1, None)
<itertools.islice object at 0x00AF9630>
>>> u''.join(_)
u'bcdefghijklmn'

ある要素より後ろの要素を取り出す

これらを組み合わせると、「ある要素より後ろの要素を得る」ジェネレータを返す ilater 関数を次のようにすっきりと書くことができます。

from itertools import islice, dropwhile
def ilater(find, iterable):
    it = dropwhile(lambda item: item != find, iterable)
    return islice(it, 1, None)

では、使用してみます。

>>> it = ilater('Mick Ronson', L)
>>> list(it)
['Ian Hunter', 'Morgan Fisher', 'Brian Ferry', 'Brian Eno', 'Phil Manzanera', 'A
ndy Mackay']

これで「ある要素以降のリスト」を得ることができました。

続いて、この islice の入出力を別のものに差し替えることができるかどうか見てみます。

出力をいじる

自作した ilater も dropwhile, islice 同様、これを実行した時点では処理は行われていません。そして、その後の処理しだいでさまざまな出力を得ることができます。

>>> u', '.join(ilater('Mick Ronson', L))
u'Ian Hunter, Morgan Fisher, Brian Ferry, Brian Eno, Phil Manzanera, Andy Mackay
'
>>> list(item.split() for item in ilater('Mick Ronson', L))
[['Ian', 'Hunter'], ['Morgan', 'Fisher'], ['Brian', 'Ferry'], ['Brian', 'Eno'],
['Phil', 'Manzanera'], ['Andy', 'Mackay']]

もうすこし複雑なことをしたい場合は for 文をまわすとよいでしょう。

>>> for name in ilater('Mick Ronson', L):
...   first, last = name.split()
...   print ''.join([first[0], '.', last[0], '.'])
...
I.H.
M.F.
B.F.
B.E.
P.M.
A.M.

入力をいじる

ilater 関数への入力はイテレート可能であれば十分です。たとえばユニコード文字列。

>>> list(ilater(u'k', u'abcdefghijklmn'))
[u'l', u'm', u'n']

遅延評価なので無限長のイテレータ相手でも破綻しません。

>>> from itertools import count
>>> it = ilater(3, count())
>>> next(it)
4
>>> next(it)
5
>>> # next するたびにいくつでも出てくる

list(itertools.count()) のようなコードを書いてしまうと処理が止まらないので、リストを作り終えてから [1:] とすることはできません。しかし、内部で islice を使用している ilater ならば、まるで [1:] したかのように先頭を取り除くことができます。 count() は 0, 1, 2 … という値を返すのですが、 0, 1, 2 は dropwhile によって、 3 は islice によって捨てられるので、 ilater(3, count()) は 4, 5, 6 … を返すジェネレータとなります*1。

結論

itertools の dropwhile, islice を用いることで、さまざまな入出力の形式に柔軟に対応できる「ある要素より後ろの要素を取り出す」コードを書くことができます。

おまけ

これで目的の達成はしたのですが、もう少し続けてみます。

オリジナルに近づける

id:mohayonao さんの later とは「戻り値がジェネレータではなくリスト」、「要素が見つからなかった時、なにも返さないのではなく、引数そのものを返す」、「第１引数と第２引数が逆」という差異があります。というわけで、 ilater を使って元の later の動作と変わらない later2 関数を作ってみます。

# オリジナル
def later(L, find):
    if L.count(find):
        L = L[::-1]
        L = L[:L.index(find)]
        return L[::-1]
    else: return L

def later2(L, find):
    u"""ilater 使用版 later"""

    result = list(ilater(find, L))
    return result if result else L

本体 2 行で書くことができました。 ilater も本体 2 行なので元と比べてもそれほど長くはなっていません。

また、検索対象が count, index メソッドを持っていなくても動作します。

複雑になっても大丈夫

Release が 1965 以前のものを取り除きます。これは自作した ilater ではなく dropwhile そのものを使うことで可能です。

>>> L = [
...     { "Title":"Please Please Me",   "Release":1963, "Peak":1 },
...     { "Title":"With The Beatles",   "Release":1963, "Peak":1 },
...     { "Title":"A Hard Day's Night", "Release":1964, "Peak":1 },
...     { "Title":"Beatles For Sale",   "Release":1964, "Peak":1 },
...     { "Title":"HELP!",       "Release":1965, "Peak":1 },
...     { "Title":"Rubber Soul", "Release":1965, "Peak":1 },
...     { "Title":"Revolver",    "Release":1966, "Peak":1 },
...     { "Title":"Sgt. Pepper's Lonely Hearts Club Band", "Release":1967, "Peak
":1 },
...     { "Title":"Magical Mystery Tour", "Release":1967, "Peak":1 },
...     { "Title":"The Beatles",      "Release":1968, "Peak":1 },
...     { "Title":"Yellow Submarine", "Release":1969, "Peak":1 },
...     { "Title":"Abbey Road", "Release":1969, "Peak":1 },
...     { "Title":"Let It Be",  "Release":1970, "Peak":1 }]
>>>
>>> it = dropwhile(lambda item: item['Release'] <= 1965, L)
>>>
>>> list(it)
[{'Release': 1966, 'Peak': 1, 'Title': 'Revolver'}, {'Release': 1967, 'Peak': 1,
 'Title': "Sgt. Pepper's Lonely Hearts Club Band"}, {'Release': 1967, 'Peak': 1,
 'Title': 'Magical Mystery Tour'}, {'Release': 1968, 'Peak': 1, 'Title': 'The Be
atles'}, {'Release': 1969, 'Peak': 1, 'Title': 'Yellow Submarine'}, {'Release':
1969, 'Peak': 1, 'Title': 'Abbey Road'}, {'Release': 1970, 'Peak': 1, 'Title': '
Let It Be'}]
>>>

汎用性と速度のトレードオフ

対象がリストやタプルと断定できるなら、つまり index メソッドが使えることが確実ならば、コメントにもあったとおり

>>> L[L.index('Mick Ronson')+1:]
['Ian Hunter', 'Morgan Fisher', 'Brian Ferry', 'Brian Eno', 'Phil Manzanera', 'A
ndy Mackay']

という index とスライスのあわせ技が速いです。対象が文字列ならば

>>> a = 'abcabcabc'
>>> a.partition('ab')[2]
'cabcabc'

などでしょうか。

たかだか一つを捨てるためだけに islice はやりすぎか？

islice を使わない方法もあります。単に値を一つ読み捨てるだけです。

>>> from itertools import dropwhile
>>> it = dropwhile(lambda item: item != 'Mick Ronson', L)
>>> next(it)
'Mick Ronson'
>>> list(it)
['Ian Hunter', 'Morgan Fisher', 'Brian Ferry', 'Brian Eno', 'Phil Manzanera', 'A
ndy Mackay']

関数化するともう少し実用的になります。

>>> def ilater2(find, iterable):
...   it = dropwhile(lambda x: x != find, iterable)
...   next(it)
...   return it
...
>>> list(ilater2('Mick Ronson', L))
['Ian Hunter', 'Morgan Fisher', 'Mick Ronson', 'Brian Ferry', 'Brian Eno', 'Phil
 Manzanera', 'Andy Mackay']

*1:4, 5, 6 ... という無限長リストが欲しいときの最適な方法は itertools.count(4) 。 4, 5, 6 … n という有限の場合は xrange(4, n + 1) 。

銀月の符号

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