dWOut = np.zeros((X.shape[1], flag.shape[1]))
idx1 = np.argwhere(flag.T > 0)
X_ext1 = X[idx1[:, (1)]]
shift_idx1 = np.append(0, np.where(np.diff(idx1[:, (0)]) > 0)[0] + 1)
dWOut[:, (np.unique(idx1[:, (0)]))] += np.add.reduceat(X_ext1, shift_idx1,
    axis=0).T
idx2 = idx1[idx1[:, (1)].argsort()]
cols_idx1 = y[idx2[:, (1)]]
X_ext2 = X[idx2[:, (1)]]
sort_idx = y[idx2[:, (1)]].argsort()
X_ext2 = X_ext2[sort_idx]
shift_idx2 = np.append(0, np.where(np.diff(cols_idx1[sort_idx]) > 0)[0] + 1)
dWOut[:, (np.unique(cols_idx1))] -= np.add.reduceat(X_ext2, shift_idx2, axis=0
    ).T
Как я могу векторизовать этот цикл for в numpy?