Python：文字列から同じ文字が連続する位置と長さを得る

たとえば、aabbcccdeeee という文字列から、a が0文字目から2文字、b が2文字目から２文字……という情報が欲しい（最初の文字を0文字目とする）。

最初に書いた関数がこれ。

def charpos(s):
    current = None
    result = []
    for i, c in enumerate(s):
        if current is None:
            current = c
            result.append([c, i, 1])
        elif c == current:
            result[-1][2] += 1
        else:
            current = c
            result.append([c, i, 1])
    return result

「文字、位置、長さが入ったリスト」のリストが返ってくる。実行してみるとこうなる。

>>> charpos('aabbcccdeeee')
[['a', 0, 2], ['b', 2, 2], ['c', 4, 3], ['d', 7, 1], ['e', 8, 4]]

期待するものはできた。けど、もっとスマートにいかないものか。分岐の current is None と else のコードが重複してるのを整理すればすっきりはするけど、素朴な実装には変わりがない（素朴なのが悪いわけではないとも思うけども）。

あと余談だけど、こういう場合には、外側のリストの要素はリストじゃなくてタプルのほうがいいと思うんだよね。

さて、itertools にある groupby を使ってみた。この関数は iterable なオブジェクトから連続する要素をグループにまとめてくれる（タプルになる）。

>>> import itertools
>>> for g in itertools.groupby('aabbcccdeeee'):
...     print(g)
...
('a', <itertools._grouper object at 0x00000207D778F760>)
('b', <itertools._grouper object at 0x00000207D778F730>)
('c', <itertools._grouper object at 0x00000207D778F760>)
('d', <itertools._grouper object at 0x00000207D778F730>)
('e', <itertools._grouper object at 0x00000207D778F760>)

各タプルの2つ目の要素は itertools._grouper のオブジェクトだけど、リストにしてやると次のようになる。

>>> for g in itertools.groupby('aabbcccdeeee'):
...     print(g[0], list(g[1]))
...
a ['a', 'a']
b ['b', 'b']
c ['c', 'c', 'c']
d ['d']
e ['e', 'e', 'e', 'e']

あとはこれを数えてやればいい。

def charpos2(s):
    idx = 0
    result = []
    for g in itertools.groupby(s):
        l = len(list(g[1]))
        result.append((g[0], idx, l))
        idx += l
    return result

>>> charpos2('aabbcccdeeee')
[('a', 0, 2), ('b', 2, 2), ('c', 4, 3), ('d', 7, 1), ('e', 8, 4)]

まぁ満足。だけどもう少し……というわけで、畳み込み関数を使ってみたらどうだろう。Haskell には mapAccumL っていう関数があるけど、Python だと itertools.accumulate が似た感じだ。こうなった。

def charpos3(s):
    g = itertools.groupby(s)
    acc = itertools.accumulate(g, lambda a, b: (b[0], a[1]+a[2], len(list(b[1]))), initial=('', 0, 0))
    return list(acc)[1:]

>>> charpos3('aabbcccdeeee')
[('a', 0, 2), ('b', 2, 2), ('c', 4, 3), ('d', 7, 1), ('e', 8, 4)]

期待通りの結果は得られたけど、これはかえって解りにくいか。