Watersheds

_watersheds.cs

using System;
using System.Collections.Generic;
using System.Linq;

List<(int x, int y)>[] FindWatersheds(int[,] map) {
    var ans = new Dictionary<int, List<(int x, int y)>>();
    var h = map.GetLength(0);
    var w = map.GetLength(1);
    var uf = new UF(w * h);
    for (int x = 0; x < w; x++) {
        for (int y = 0; y < h; y++) {
            // find min
            var (minX, minY) = FindMin(map, x, y, w, h);
            // connect components
            uf.merge(y * w + x, minY * w + minX);
        }
    }
    // retrieve disjoint sets
    for (int x = 0; x < w; x++) {
        for (int y = 0; y < h; y++) {
            var watershedId = uf.find(y * w + x);
            (ans[watershedId] = ans.GetValueOrDefault(watershedId, new List<(int x, int y)>()))
                .Add((x, y));
        }
    }
    return ans.Values.ToArray();
}

(int x, int y) FindMin(int[,] map, int x, int y, int w, int h) => Neighbors(x, y)
    .Where(p => p.x >= 0 && p.x < w && p.y >= 0 && p.y < h)
    .Select(p => (map[p.y, p.x], p))
    .Min()
    .Item2;

IEnumerable<(int x, int y)> Neighbors(int x, int y) {
    for (int xx = -1; xx <= 1; xx++)
        for (int yy = -1; yy <= 1; yy++)
            yield return (x + xx, y + yy);
}

int[,] input = {
    {0, 1, 2, 3, 2, 1},
    {1, 1, 2, 3, 2, 0},
    {2, 2, 3, 3, 2, 1},
};
var watersheds = FindWatersheds(input);
var id = 1;
foreach (var watershed in watersheds) {
    foreach (var (x, y) in watershed) {
        input[y, x] = id;
    }
    id++;
}

for (int y = 0; y < input.GetLength(0); y++) {
    for (int x = 0; x < input.GetLength(1); x++) {
        Console.Write(input[y, x] + " ");
    }
    Console.WriteLine();
}

uf.cs

    public class UF
    {
        public int[] id { get; set; }
        public int[] sz { get; set; }
        public int cnt { get; set; }

        public UF(int N)
        {
            cnt = N;

            id = new int[N];
            sz = new int[N];

            for (int i = 0; i < N; i++)
            {
                id[i] = i;
                sz[i] = 1;
            }
        }

        // Return the id of component corresponding to object p.
        public int find(int p)
        {

            int root = p;
            while (root != id[root])
                root = id[root];

            while (p != root)
            {
                int newp = id[p];
                id[p] = root;
                p = newp;
            }

            return root;
        }

        public void merge(int x, int y)
        {
            int i = find(x);
            int j = find(y);
            if (i == j) return;

            if (sz[i] < sz[j])
            {
                id[i] = j;
                sz[j] += sz[i];
            }
            else
            {
                id[j] = i;
                sz[i] += sz[j];
            }

            cnt--;
        }


        public bool connected(int x, int y)
        {
            return find(x) == find(y);
        }

        public int count()
        {
            return cnt;
        }
    }