Pythonで多次元リストの全要素数をカウントする方法と速さ検証

何かの機会で2次元リストの全要素数をカウントしたくなった(何の機会でしたくなったのかは忘れた)のでその方法をまとめます．

2次元リスト

このようなリスト作りました．このリストの全要素数をカウントします．

d2_list = [[i,i+1] for i in range(100)] + [[i, i*2, i*3] for i in range(100)] + [[i, i+2, i+4, i+5] for i in range(100)]

内包表記を使う

sum(len(i) for i in d2_list)
>>900

sumの中の内包表記で各リストの要素数を計算しsumでその和を取ることで全要素をカウントしています．

sum(list, [])を使う

len(sum(d2_list, []))
>>900

sum(list, [])は2次元リストを1次元にしてくれます．

a = [[1, 2], [3, 4], [5, 6]]
sum(a, [])
>>[1, 2, 3, 4, 5, 6]

1次元に変換してその長さをとれば全要素数をカウントできます．

ndarrayに変換して.sizeを取る

ndarrayは多次元配列の時全てのリストの長さが同じで無いといけません．なのでd2_listではこの方法は使えません

d2_arr = np.array([[1, 2], [3, 4], [5, 6]])
d2_arr.size
>>6

多次元配列の時

ndarrayに変換して.sizeを取る

2次元の時と同じく全てのリストの長さが同じで無いといけませんが，同じなら多次元でもカウントできます．

d2_arr = np.array([[[1, 2], [3, 4], [5, 6]], [[1, 2], [3, 4], [5, 6]]])
d2_arr.size
>>12

再帰を使う

def multid_len(l):
    count = 0
    if isinstance(l, list):
        for v in l:
            count += multid_len(v)
        return count
    else:
        return 1
md_list = [[[i, j*2] for i in range(100)] for j in range(100)] + [[i, i**2] for i in range(100)] + [1, 2, 3]
multid_len(md_list)
>>20203

isinstanceを使うと型が等しいかチェックしてくれます．lがリストだった場合，各要素に対しmultid_lenを再帰的に実行します．lがリストでなかった場合それが要素なので1を返しcountに足されます．

他に多次元配列の全要素数を数える方法あるでしょうか……？あればコメントで教えて下さい．私は思い付きませんでした．

速さ比べ

2次元配列かつ要素数が全て同じ

import time
import numpy as np
d2_list = [[i,i+1] for i in range(10**6)]
end_times = 0
for i in range(10):
  start = time.time()
  print(d2_len(d2_list))
  end_time = time.time() - start
  print(end_time)
  end_times += end_time
print("10回平均", end_times/10)

このような実験コードを作りました．速度比較は初めてやるので不十分な所があるかもしれません．
d2_lenにnumpy, 内包表記, sumのそれぞれを書き，結果を確認します．実際はこのコードを更に10回実行し平均を取りました．環境はgoogle colabです．

	時間
numpy	0.9028692404429117
内包表記	0.10396323204040529
sum	30秒以上

sumがこんなに遅いとは思いませんでした．内包表記は早そうだなと思ってましたがこんなに差がつくとは……
numpyより内包表記が9倍も早いのも意外でした．d2_lenの中にnp.array()が含まれているので，あらかじめndarrayに変換してから実行した場合 0.0003843545913696289秒まで早くなりました．
やはりndarrayに変換するのに時間が掛かってるようです．全体としてndarrayで処理するときはsizeを，リストの処理だけで済むときは内包表記を使うのが良いと分かりました．
要素数がバラバラの時の内包表記とsumも比べようかと思いましたが結果は目に見えてるのでやりません．

結論

ndarrayで処理したいならsize, 使わずに処理できるのならsum(len(i) for i in list)を使う．