How to perform outer subtraction along an axis in numpy-CodePudding

I used to perform an outer subtraction on two one-dimensional arrays as follows to receive a single two-dimensional arrays that contains all pairs of subtractions:

import numpy as np

a = np.arange(5)
b = np.arange(3)
result = np.subtract.outer(a, b)
assert result.shape == (5, 3)
assert np.all(result == np.array([[aa - bb for bb in b] for aa in a ])) # no rounding errors

Now the state space switches to two dimensions, and I would like to perform the same operation, but only perform each subtraction on the two values on the last axis of the arrays A and B:

import numpy as np

A = np.arange(5 * 2).reshape(-1, 2)
B = np.arange(3 * 2).reshape(-1, 2)
result = np.subtract.outer(A, B)

# Obviously the following does not hold, because here we have got all subtractions, therefore the shape (5, 2, 3, 2)
# I would like to exchange np.subtract.outer such that the following holds:
# assert result.shape == (5, 3, 2)

expected_result = np.array([[aa - bb for bb in B] for aa in A ])
assert expected_result.shape == (5, 3, 2)

# That's what I want to hold:
# assert np.all(result == expected_result) # no rounding errors

Is there a "numpy-only" solution to perform this operation?

CodePudding user response：

You can expand/reshape A to (5, 1, 2) and B to (1, 3, 2) and let the broadcasting do the job:

A[:, None, :] - B[None, :, :]

CodePudding user response：

A[:, None] - B[None, :] does it.

A = np.arange(5 * 2).reshape(-1, 2)
B = np.arange(3 * 2).reshape(-1, 2)

expected_result = np.array([[aa - bb for bb in B] for aa in A ])

C = A[:, None] - B[None, :]

np.allclose(expected_result, C)
#> True

The exact same syntax works for your first example too. This is because with your requirement, you are combining every first axis element of A with every first axis element of B.