1D implementation of Hoshen-Kopelman cluster labelling algorithm to identify different trips in a e-scooters data frame

hoshen_kopelman_1d.py

import random
import pandas as pd
import numpy as np

lock_state = ["locked", "unlocked", "unlocked", "unlocked", "locked", "locked", "unlocked", "unlocked"]

random_values = random.sample(range(2,20), 8) 

df = pd.DataFrame(
    {'state': lock_state,
     'random': random_values
    })
    

def hoshen_kopelman_1d(grid, occupied_label):
    """
    Hoshen Kopelman implementation for 1D graphs.
    
    Parameters:
            grid (pd.DataFrame): The 1D grid. 
            ocuppied_label (str): the label that identifies occupied nodes.

    Returns:
            labeled_grid (pd.DataFrame): grid with cluster labeled nodes.
    """
    
    # create labeled_grid and assign all nodes to cluster 0
    labeled_grid = df.assign(cluster=0)
    cluster_count = 0
    
    # iterate through the grid's nodes left to right
    for index, node in grid.iterrows():
        # check if node is occupied
        if node["state"] == occupied_label: # node is occupied
            if index == 0:
                # initialize new cluster
                cluster_count += 1
                labeled_grid.loc[0, "cluster"] = cluster_count
            else:
                # check if left-neighbour node is occupied
                if labeled_grid.loc[index-1, "cluster"] != 0: # left-neighbour node is occupied
                    # assign node to the same cluster as left-neighbour node
                    labeled_grid.loc[index, "cluster"] = labeled_grid.loc[index-1, "cluster"]
                else: # left-neighbour node is unoccupied
                    # initialize new cluster
                    cluster_count += 1
                    labeled_grid.loc[index, "cluster"] = cluster_count
                    
    return labeled_grid
                

M = hoshen_kopelman_1d(grid=df, occupied_label="unlocked")

# to retrieve trip_1
trip_1 = M.loc[M['cluster'] == 1]