zipwithとunzip

zipwithを作ってみました。zipwithは複数のシーケンスを受け取って、それぞれの要素に任意の関数を適用させた結果を要素とするリストを返す関数です。

def zipwith(f, *ls):
  return map(f, *ls)

l1 = range(10)
l2 = [i%3 for i in range(10)]

print zipwith(lambda x, y: x+y, l1, l2)
# => [0, 2, 4, 3, 5, 7, 6, 8, 10, 9]

print zipwith(lambda x, y, z: x*y*z, l1, l2, range(8))
# => TypeError: unsupported operand type(s) for *: 'int' and 'NoneType'

わざわざdefするまでもなかったかな……などと思っていたのですが、長さが違うリストを渡したらエラーになってしまいました。改良してみます。

def zipwith(f, *ls):
  return [f(*i) for i in zip(*ls)]

l1 = range(10)
l2 = [i%3 for i in range(10)]

print zipwith(lambda x, y: x+y, l1, l2)
# => [0, 2, 4, 3, 5, 7, 6, 8, 10, 9]

print zipwith(lambda x, y, z: x*y*z, l1, l2, range(8))
# => [0, 1, 8, 0, 16, 50, 0, 49]

ついでにunzipも作ってみました。unzipはzipでまとめたリストを引数として受け取って、元のリストに戻す関数です。

def unzip(l):
  r = [[] for i in l[0]]
  for i in l:
    for j in range(len(r)):
      r[j].append(i[j])
  return r

l1 = range(5)
l2 = [i**2 for i in range(5)]

z = zip(l1, l2)
print z
# => [(0, 0), (1, 1), (2, 4), (3, 9), (4, 16)]

print unzip(z)
# => [[0, 1, 2, 3, 4], [0, 1, 4, 9, 16]]

z = zip(range(10), [i**2 for i in range(8)], 'apple juice')
print z
# => [(0, 0, 'a'), (1, 1, 'p'), (2, 4, 'p'), (3, 9, 'l'), (4, 16, 'e'), (5, 25, ' '), (6, 36, 'j'), (7, 49, 'u')]

print unzip(z)
# => [[0, 1, 2, 3, 4, 5, 6, 7], [0, 1, 4, 9, 16, 25, 36, 49], ['a', 'p', 'p', 'l', 'e', ' ', 'j', 'u']]

うーん、もっと簡潔に書けそうな気がする……。