def _get_leaf_node(vec: Dict, assignment: np.ndarray) -> np.ndarray:
    """
    Get the sequence number of the selected leaf node from
    the leaf node selection vector
    Args:
        vec: vectorized information dict in a tree, refer to
             _vectorize_tree(tree)
        assignment: one hot selection vector of leaf, shape:
                    (N, leaf_node_num),N is number of data,
                    leaf_node_num is the node num in a tree.

    Returns:
        (N,) The sequence number of the leaf node selected by all data
    """
    is_leaf = vec['is_leaf']
    leaf = np.arange(is_leaf.size)[is_leaf]
    return (assignment * leaf).sum(1).astype(np.int32)

    Get the sequence number of the selected leaf node from
    the leaf node selection vector
    Args:
        vec: vectorized information dict in a tree, refer to
             _vectorize_tree(tree)
        assignment: one hot selection vector of leaf, shape:
                    (N, leaf_node_num),N is number of data,
                    leaf_node_num is the node num in a tree.

    Returns:
        (N,) The sequence number of the leaf node selected by all data
    
fedlearner/model/tree/tree.py