Behzad Jamali
Behzad Jamali

Reputation: 924

Grouping adjacent equal elements of a 2D numpy array

I wonder if there is a function to group equal, connected elements of a 2D array like

a = np.array([[12,12,14,14,11,11],
              [10,10,11,11,11,11],
              [10,14,14,10,11,13],
              [12,12,14,13,13,13]])

into an array this:

            [[1, 1, 2 ,2, 3, 3],
             [4, 4, 3, 3, 3, 3],
             [4, 5, 5, 6, 3 ,7],
             [8, 8, 5, 7, 7, 7]]

The rules for connection: an element [i, j] is connected to [i-1, j], [i+1, j], [i, j-1], and [i, j+1].

I found scipy.ndimage.measurements.label but the problem is that it just consider the array values as zero (background) and ones.

Upvotes: 1

Views: 1032

Answers (1)

user6655984
user6655984

Reputation:

Depending on the number of unique values, it may be practical to simply use label in a loop, adding the results with appropriate offsets. The offsets are needed because after, say, the first 3 features are labeled, the labels for subsequent ones should begin with 4, and so on.

from scipy.ndimage import label
values = np.unique(a.ravel())
offset = 0
result = np.zeros_like(a)
for v in values:
  labeled, num_features = label(a == v)
  result += labeled + offset*(labeled > 0)
  offset += num_features
print(result)

This prints

[[4 4 7 7 3 3]
 [1 1 3 3 3 3]
 [1 8 8 2 3 6]
 [5 5 8 6 6 6]]

which is the same as your expected result, up to permutation of labels (which don't have any meaning anyway).

Upvotes: 2

Related Questions