numpy の複数回インクリメント

Question

ある長さ N の np.array A と、重複を許容する index たちの集合 I があります。
I の各要素に対して同時にインクリメントなどの操作を行う実装は numpy にありますでしょうか？

具体例は以下で、実際に取得したいのは [4 2 4 5 5] となります。

import numpy as np

A = np.array([1, 2, 3, 4, 5])
I = [0, 2, 0, 3, 0]
A[I] += 1
print(A)  # [2 2 4 5 5]

I を iterate するしかないのでしょうか？ご教授いただけると幸いです。
よろしくお願い致します。

Answer 1

解決済みとなっていますが、より容易な方法があるので紹介します。

np.ufunc.at()で特定のインデックスに対して複数回ユニバーサル関数を適用させることができます。
足し算を実行するユニバーサル関数np.add()を用いる場合はnp.add.at()の形となります。

import numpy as np

A = np.array([1, 2, 3, 4, 5])
I = [0, 2, 0, 3, 0]

np.add.at(A, I, 1)
print(A)  # [4 2 4 5 5]

Comments

[moni]

ありがとうございます！
ufunc にこんな機能もあったのですね、まだまだ調べ足りないのでドキュメントを見に行こうと思います。

Answer 2

numpy.bincount()を使ってIの中の各indexの個数を求めて、それを足してやるのはどうでしょうか。

import numpy as np

A = np.array([1, 2, 3, 4, 5])
I = [0, 2, 0, 3, 0]

count = np.bincount(I, minlength=len(A))
print(count)
# [3 0 1 1 0]

A += count
print(A)
# [4 2 4 5 5]

Comments

[moni]

まさしく欲しかったものです！bincount という関数があるのを初めて知りました、ありがとうございます。

Answer 3

import numpy as np

A = np.array([1, 2, 3, 4, 5])
I = [0, 2, 0, 3, 0]

u, c = np.unique(I, return_counts=True)
print(u) # [0 2 3]
print(c) # [3 1 1]

A[u] += c

print(A) # [4 2 4 5 5]

Comments

[moni]

unique を取る際に counts も同時に行えるのですね、知りませんでした。ありがとうございます！