Compute Squared Distances to Reference Rows

Task

Write compute_squared_distances(queries, references).

Each row of references is one reference point. references is a non-empty two-dimensional finite numeric NumPy array with shape (reference_count, feature_count). queries is either one finite numeric row with shape (feature_count,) or a non-empty two-dimensional array with shape (query_count, feature_count). The feature counts must agree. Invalid input raises ValueError with "invalid distance arrays".

For a query q and reference row r, calculate the squared distance by the fully stated rule

distance(q, r) = sum((q[j] - r[j]) ** 2 for every feature j)

Do not take a square root. Return a new floating-point NumPy array. A one-dimensional query returns shape (reference_count,); a two-dimensional query returns shape (query_count, reference_count). The result is ordered by query first and reference row second. Neither input may change.

Example

The result is [4, 4, 2]: for example, (1 - 0) ** 2 + (2 - 3) ** 2 == 2 for the last row. With queries [[1, 2], [0, 1]], the result has shape (2, 3) and its second row is [1, 10, 5].

Your implementation

Edit solution.py and keep this function signature:

You may import NumPy as np. Compute in a type that returns floating-point distances even when the inputs contain integers. Do not modify an input, print, or ask for input.