Is there a way to vectorize cumulatively decrementing duplicate values?

[ahuang11]

Goal: for each row, find duplicates and then for the first duplicate, keep as is, second duplicate, subtract 1, third duplicate subtract 2, and so on.

Ideally, it would be nice to find a way to vectorize this.

import numpy as np
import xarray as xr

def offset_duplicate(da_sel):
    unique_vals = np.unique(da_sel)
    for val in unique_vals:
        index = da_sel == val
        if index.sum() == 1:
            continue
        da_sel = da_sel.where(~index, da_sel - index.cumsum() + 1)
    return da_sel

da = xr.DataArray(np.array([[0, 3, 3, 3, 4], [0, 1, 4, 4, 4]]), dims=("row", "col"))
da_want = xr.DataArray(np.array([[0, 3, 2, 1, 4], [0, 1, 4, 3, 2]]), dims=("row", "col"))

da_offset = xr.concat([offset_duplicate(da.sel(row=row)) for row in da["row"]], "row")

[ahuang11]

Or framed another way, I don't want the ranks to be the same:

import xarray as xr
import numpy as np
da = xr.DataArray(np.array([[0, 3, 3, 3, 4], [0, 1, 4, 4, 4]]), dims=("row", "col"))
da.rank("col")

https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.rank.html
I want first: ranks assigned in order they appear in the array

[ahuang11]

Oh, as long as it's 2D, I can transform xarray into pandas and back to xarray

import xarray as xr
import pandas as pd
da = xr.DataArray(np.array([[6, 3, 3, 3, 4], [0, 1, 4, 4, 4]]), dims=("row", "col"))

df = pd.DataFrame(da.values)
xr.DataArray(df.rank(method="first", axis=1), dims=da.dims)