ultimate_beneficial_owner.py

from dataclasses import dataclass
from typing import Tuple, Type, Dict

VertexType = Type['Vertex']


@dataclass
class Vertex:
    name: str
    value: int
    parent: VertexType = None
    is_ubo: bool = False
    children: Tuple[VertexType] = tuple()


def get_beneficial_owners_example2() -> Vertex:
    '''
    Example 2:
    - Input (beneficial owners):
            C
            /      \
        C1(50%)  I1(50%)
        /    \
    I2(50%)  I3(50%)

    - Output (ultimate beneficial owners):
    I1 (50%), I2 (25%), I3 (25%)
    '''

    C = Vertex(name='C', value=100)
    C1 = Vertex(name='C1', value=50, parent=C)
    I1 = Vertex(name='I1', value=50, parent=C)

    I2 = Vertex(name='I2', value=50, parent=C1)
    I3 = Vertex(name='I3', value=50, parent=C1)

    C.children = (C1, I1)
    C1.children = (I2, I3)
    return C


def get_beneficial_owners_example3() -> Vertex:
    '''
    Example 3:
    - Input (beneficial owners):
                C
                /      \
            C1(50%)  I1(50%)
            /    \
        C2(40%)  I1(60%)
        /     \
    I2(10%) I3(90%)

    - Output (ultimate beneficial owners):
    I1 (80%), I2 (2%), I3 (18%)
    '''

    C = Vertex(name='C', value=100)

    C1 = Vertex(name='C1', value=50, parent=C)
    I1 = Vertex(name='I1', value=50, parent=C)

    C2 = Vertex(name='C2', value=40, parent=C1)
    I1_C1 = Vertex(name='I1', value=60, parent=C1)

    I2 = Vertex(name='I2', value=10, parent=C2)
    I3 = Vertex(name='I3', value=90, parent=C2)

    C.children = (C1, I1)
    C1.children = (C2, I1_C1)
    C2.children = (I2, I3)
    return C


def compute_val(vx: Vertex) -> int:
    if vx.parent:
        return (vx.value / 100) * compute_val(vx.parent)
    else:
        return vx.value


def walk_graph(parent: Vertex, ubo_values: Dict = dict()) -> dict:
    for c in parent.children:
        if c.children:
            walk_graph(c, ubo_values)
        else:
            c.is_ubo = True
            ubo_values[c.name] = compute_val(c) + ubo_values.get(c.name, 0)
    return ubo_values


print(walk_graph(get_beneficial_owners_example2(), dict()))
# {'I2': 25.0, 'I3': 25.0, 'I1': 50.0}

print(walk_graph(get_beneficial_owners_example3(), dict()))
# {'I2': 2.0, 'I3': 18.0, 'I1': 80.0}