pythonでの配列の扱い方についての質問です。

Question

質問文がふわっとしていてすみません。どう書けばいいのか分かりませんでした。

前提・実現したいこと

a=[1,2,3,4,5,...,99,100]
b=[1,16,31,76]
という2つの配列があったときに、
「配列bの各要素から始まって、5つ分の要素をaから取り出す」という操作をしたいです。
つまり、
c=[1,2,3,4,5, 16,17,18,19,20, 31,32,33,34,35, 76,77,78,79,80]
という配列を作りたいのですが、スマートな書き方を教えていただきたいです。

今は

c= []
for num in b:
    c += [a[num], a[num+1], a[num+2], a[num+3], a[num+4]]
c = np.array(c)

という書き方をしています。これでも動くのですが何となくカッコ悪いなあと感じています。
内包表記とか使うとスマートにかけるという噂は聞いているのですが、お恥ずかしながらまだ全然使いこなせていません。

もしスマートに書ける方法がございましたら教えていただけますと嬉しいです。

Answer 1

内包表記を使ってみます。

>>> [e for start in b for e in a[start-1:start+4]]
[1, 2, 3, 4, 5, 16, 17, 18, 19, 20, 31, 32, 33, 34, 35, 76, 77, 78, 79, 80]

慣れていないとなかなか読みづらいので、スライスの活用に留めても良いでしょう。

>>> c = []
>>> for start in b:
...     c += a[start-1:start+4]
...
>>> c
[1, 2, 3, 4, 5, 16, 17, 18, 19, 20, 31, 32, 33, 34, 35, 76, 77, 78, 79, 80]

NumPyでもっとサクッと書く方法もあるかもしれませんが。

Comments

[siruku6]

リスト内包表記をうまく改行すると、それなりに見やすくはなると聞いたことがあります。
処理速度は変わらないので、内包表記を if や for の位置などで改行するのは行うことがあります。

[LouiS0616]

そうですね。行が長くなるようなら適宜改行をした方が良いかと思います。

[taku_t]

おおお！ありがとうございます！
内包表記の中に内包表記が入ってますね...。ちょっと考えてみます>_<
そしてスライス操作の方もありがとうございます！言われてみれば...って感じです。もっと見やすくシンプルなコードを書けるように精進します。

[taku_t]

と思ったらスライス操作はエラーでした。
cの[0,]という型と、a[start-1:start+4]の[5,]という型がかみ合ってないようです。。もう少し考えてみます。

Answer 2

必ずしもスマートなコードでなくて良いと思います。
コードの読みやすさや、後で自身で容易に修正出来ることも重要です。