@JavierRuano I find the simpler solution from a similar question in stack overflow.

sort_pair = np.take_along_axis(pair.values, cld.argsort(axis=0), axis=0)

Complete example

``` import xarray as xr import numpy as np

x = 4 y = 2 z = 4

data = np.arange(xyz).reshape(z, y, x)

3d array with coords

cld_1 = xr.DataArray(data, dims=['z', 'y', 'x'], coords={'z': np.arange(z)})

2d array without coords

cld_2 = xr.DataArray(np.arange(xy).reshape(y, x)1.5+1, dims=['y', 'x'])

expand 2d to 3d

cld_2 = cld_2.expand_dims(z=[4])

concat

cld = xr.concat([cld_1, cld_2], dim='z')

paired array

pair = cld.copy(data=np.arange(xy(z+1)).reshape(z+1, y, x))

sort_pair = np.take_along_axis(pair.values, cld.argsort(axis=0), axis=0)

print(cld) print(pair) print(sort_pair) ```

Output: ``` <xarray.DataArray (z: 5, y: 2, x: 4)> array([[[ 0. , 1. , 2. , 3. ], [ 4. , 5. , 6. , 7. ]],

   [[ 8. ,  9. , 10. , 11. ],
    [12. , 13. , 14. , 15. ]],

   [[16. , 17. , 18. , 19. ],
    [20. , 21. , 22. , 23. ]],

   [[24. , 25. , 26. , 27. ],
    [28. , 29. , 30. , 31. ]],

   [[ 1. ,  2.5,  4. ,  5.5],
    [ 7. ,  8.5, 10. , 11.5]]])

Coordinates: * z (z) int64 0 1 2 3 4 Dimensions without coordinates: y, x <xarray.DataArray (z: 5, y: 2, x: 4)> array([[[ 0, 1, 2, 3], [ 4, 5, 6, 7]],

   [[ 8,  9, 10, 11],
    [12, 13, 14, 15]],

   [[16, 17, 18, 19],
    [20, 21, 22, 23]],

   [[24, 25, 26, 27],
    [28, 29, 30, 31]],

   [[32, 33, 34, 35],
    [36, 37, 38, 39]]])

Coordinates: * z (z) int64 0 1 2 3 4 Dimensions without coordinates: y, x [[[ 0 1 2 3] [ 4 5 6 7]]

[[32 33 34 35] [36 37 38 39]]

[[ 8 9 10 11] [12 13 14 15]]

[[16 17 18 19] [20 21 22 23]]

[[24 25 26 27] ```

Note, I have to use pair.values instead of pair in the last sorting step. Otherwise, I will get this error:

IndexError: Unlabeled multi-dimensional array cannot be used for indexing: y