pythonのL2ノルムを取るプログラムのfor文を使わない書き方, 高速化について

Question

前提・実現したいこと

(130433, 2)のnumpyの配列Aがあります。 ２次元目でL2ノルムをとった値によってAを割ることで配列Aの値を正規化するプログラムのfor文を使わない書き方を教えてほしいです。
以下のプログラムにて実行はできているのですが、for文を使っていたりと実行時間が長いです。
map関数等をつかったほうがいいと思うのですが、あまりなれておらず調べられませんでした。

どの関数を使ったほうがいい等簡単なアドバイスでも良いのでいただけるとありがたいです。

該当のソースコード

print(A.shape)
# (130433, 2)

# abs_Aに正規化した値を保存
abs_A = np.zeros((A.shape[0], A.shape[1]))
for i in range(A.shape[0]):
    abs_A = abs_A/np.linalg.norm(A, axis=1)[i]

実行環境

python3.x

Answer 1

(130433, 2) の行列の列方向のノルムを計算し、(130433,) という配列を得たあと、(130433,1) に形状を変えてから除算すれば、ブロードキャストされて希望の結果が得られるかと思います。

import numpy as np

A = np.random.rand(130433, 2)
B = A / np.linalg.norm(A, axis=1)[:, np.newaxis]

# 列方向のノルムが1になっているかどうか
assert np.allclose(np.linalg.norm(B, axis=1), 1)  # True

Comments

[Takeru0208]

素早い解答ありがとうございました！
形状を変えればブロードキャストできないというエラーも解決できるのですね、勉強になりました！