jackee777のブログ

情報系学生のつぶやき

python の転置 zip(*list) を理解する

初見では,???だったけど,書いてるうちに理解した・・・

タプル型への転置は list(zip(*L))

リスト型への転置は list(map(list, (zip(*L))))

※全て zip で囲むのは,zip object, map object となるのを避けるため.複数で受け取る場合は,m, n = map(list, (zip(*L)) とも書ける

追ってみていくと,次のような感じ python で numpy を使わずに,転置(transpose)する場合 

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
L = list(zip(*L))
L
>>> L
[(1, 4, 7), (2, 5, 8), (3, 6, 9)]

このままではタプル型なので,リスト型にする時は

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
L = list(map(list, (zip(*L))))
L
>>> L
[[1, 4, 7], [2, 5, 8], [3, 6, 9]]

ちなみに,numpy を使うならで良い

import numpy as np
L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
L = np.asarray(L).T
L
>>> L
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])

最初から numpy 使えば良いよね,と言われればそうなのだが, この記法-zip(*L)-は直感的でなく,理由が気になったので調べる.

とは言っても,zip(L) と zip(*L) の違いを把握すれば話は終わる.

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
for i in zip(L):
    print(i)

for i in zip(*L):
    print(i)

# copy しやすいようにしてるだけ
>>> for i in zip(L):
...     print(i)
...
([1, 2, 3],)
([4, 5, 6],)
([7, 8, 9],)
>>> for i in zip(*L):
...     print(i)
...
(1, 4, 7)
(2, 5, 8)
(3, 6, 9)

実行すれば,わかるが既に転置している. この実行結果を見てから気づいたのだが,よくよく考えると, zip(*L) とは zip(L[0], L[1], L[2]) と同じ処理をしていることになる. *L と書くと,L は unpack されて行ごとに分解されるため,*L = (L[0], L[1], L[2])

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
print(*L)
print((L[0], L[1], L[2]))
for i in zip(L[0], L[1], L[2]):
    print(i)

# copy しやすいようにしてるだけ
>>> for i in zip(L[0], L[1], L[2]):
...     print(i)
...
(1, 4, 7)
(2, 5, 8)
(3, 6, 9)

python 2?における実装はかなりシンプルで,初めての人もこれくらい書けそう. PEP 201 -- Lockstep Iteration | Python.org

def zip(*args):
    if not args:
        raise TypeError('zip() expects one or more sequence arguments')
    ret = []
    i = 0
    try:
        while 1:
            item = []
            for s in args:
                item.append(s[i])
            ret.append(tuple(item))
            i = i + 1
    except IndexError:
        return ret

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
for i in zip(*L):
    print(i)

# copy しやすいようにしてるだけ

python 3 では,map, zip, filter といった関数は次のように iterator 形式なっていて 無駄な try とかないし,yield によってメモリ消費量の削減を考えているぽい. Built-in Functions — Python 3.7.3 documentation

def zip(*iterables):
    # zip('ABCD', 'xy') --> Ax By
    sentinel = object()
    iterators = [iter(it) for it in iterables]
    while iterators:
        result = []
        for it in iterators:
            elem = next(it, sentinel)
            if elem is sentinel:
                return
            result.append(elem)
        yield tuple(result)

L = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
for i in zip(*L):
    print(i)

# copy しやすいようにしてるだけ

一応,python 3 でメモリ消費が少なくなっているのをチェックしました. AtCoder のコードテストのメモリ消費量によって確認しています(使い方としてはよろしくないです.すいません). こんな感じのコード

N = 100
L = [[j for j in range(N)] for i in range(N)]

# それぞれ def zip(*args) を書いて実行

for i in zip(L):
  pass

やっぱりメモリ消費量がわずかながら抑えられてましたね. なんというか,3000行3000列でやっとこさ,0.2MB というのは 普段からそこまで iterator を意識しなくても良いというのと プログラミング言語としては組み込み関数が洗練されていくのは良いって感じますね

zip の種類 N=100 N=1000 N=3000
上の zip (python 2?) 5,532 KB 37,904 KB 345,424 KB
下の zip (python 3.7.3) 5,516 KB 37,732 KB 345,200 KB

はい,以上です.