def _assign_direction_to_node(vec: Dict, direction: np.ndarray) -> np.ndarray:
    """
    According to the comparison result of the previous step,
    determine the traversal direction of each node, assign 0 or
    1 to the child nodes according to the traversal direction,
    and then set all child nodes that belong to the not-own
    parent nodes to 1
    Args:
        vec: vectorized information dict in a tree, with reference
             to response vec of function _vectorize_tree
        direction: Comparison results of features and thresholds,
                   shape: (N, node_num), N is number of data,
                   node_num is the node num in a tree.

    Returns:
        selected: assignment result for each node, shape: (N, node_num),
                    N is number of data, node_num is the node num in a tree.
    """
    N = direction.shape[0]
    node_num = vec['feature_id'].size

    col = np.tile(np.arange(node_num), (N, 1))
    selected_child = vec['children'][direction.astype(np.int32), col]
    non_leaf_num = vec['is_leaf'].size - vec['is_leaf'].sum()
    non_leaf = ~np.tile(vec['is_leaf'], (N, 1))
    node_ids = np.tile(np.arange(node_num), (N, 1))

    tmp = [np.in1d(node_ids[i], assign) for i, assign in
           enumerate(selected_child[non_leaf].reshape(N, non_leaf_num))]
    assignment = np.concatenate(tmp, axis=0).reshape(N, node_num)

    own_id = np.arange(node_num)[vec['is_owner']]
    not_own = (~np.in1d(vec['parent'], own_id))

    selected = (assignment | not_own).astype(np.int32)
    return selected

    According to the comparison result of the previous step,
    determine the traversal direction of each node, assign 0 or
    1 to the child nodes according to the traversal direction,
    and then set all child nodes that belong to the not-own
    parent nodes to 1
    Args:
        vec: vectorized information dict in a tree, with reference
             to response vec of function _vectorize_tree
        direction: Comparison results of features and thresholds,
                   shape: (N, node_num), N is number of data,
                   node_num is the node num in a tree.

    Returns:
        selected: assignment result for each node, shape: (N, node_num),
                    N is number of data, node_num is the node num in a tree.
    
fedlearner/model/tree/tree.py