Source code for wield.control.ss_bare.tarjan

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: © 2022 California Institute of Technology.
# SPDX-FileCopyrightText: © 2022 Lee McCuller <mcculler@caltech.edu>
# NOTICE: authors should document their contributions in concisely in NOTICE
# with details inline in source files, comments, and docstrings.
"""
"""
import numpy as np




def find_sccs(graph: dict[any, list[any]]):
    """
    Takes a dictionary of nodes (can be anything) that map to a list of nodes.
    The list represents the outgoing edges.

    returns: list[list[any]]

    The algorithm returns a list of lists of nodes. The outer list is topologically sorted
    and the inner list are groups of strongly-connected components.
    """
    ordnum = 0
    stack = []
    low_link = {}
    visited = set()
    sccs = []

    def _dfs(v):
        nonlocal ordnum
        v_ordnum = ordnum
        low_link[v] = ordnum
        ordnum += 1
        stack.append(v)
        visited.add(v)

        for neighbor in graph[v]:
            if neighbor not in low_link:
                _dfs(neighbor)
                low_link[v] = min(low_link[v], low_link[neighbor])
            elif neighbor in stack:
                low_link[v] = min(low_link[v], low_link[neighbor])

        if low_link[v] == v_ordnum:
            scc = []
            while True:
                u = stack.pop()
                scc.append(u)
                if u == v:
                    break
            sccs.append(scc)

    for v in graph:
        if v not in low_link:
            _dfs(v)

    return sccs





def find_sccs_matrix(matrix: np.array):
    """
    Takes a matrix. Nonzero elements represent edges.

    returns: list[list[int]]

    The algorithm returns a list of lists of columns. The outer list is topologically sorted
    and the inner list are groups of strongly-connected components of the edges given by the matrix.
    """
    assert (len(matrix.shape) == 2)
    assert (matrix.shape[0] == matrix.shape[1])
    n = matrix.shape[0]

    ordnum = 0
    stack = []
    low_link = np.empty(shape=(n,), dtype=int)
    low_link[:] = -1
    visited = np.zeros(shape=(n,), dtype=bool)
    sccs = []

    def _dfs(col):
        nonlocal ordnum
        v_ordnum = ordnum
        low_link[col] = ordnum
        ordnum += 1
        stack.append(col)
        visited[col] = True

        for row, val in enumerate(matrix[:, col]):
            if val == 0:
                continue

            if low_link[row] == -1:
                _dfs(row)
                low_link[col] = min(low_link[col], low_link[row])
            elif row in stack:
                low_link[col] = min(low_link[col], low_link[row])

        if low_link[col] == v_ordnum:
            scc = []
            while True:
                u = stack.pop()
                scc.append(u)
                if u == col:
                    break
            sccs.append(scc)

    for col in range(n):
        if low_link[col] == -1:
            _dfs(col)

    return sccs