animal_transport.py

# source: https://www.hackerrank.com/challenges/animal-transport/problem
# video: https://youtu.be/35cgDwzPIxg
inp = input
ran = range
_int = int

ref = [
    "Q3JlYXRlZCBieSBHbGl0Y2hlZCBGYWlsdXJl",
    "aHR0cHM6Ly93d3cueW91dHViZS5jb20vY2hhbm5lbC9VQ0VyU05pRFpWNHJKQ05COE5yREdSRUE="
]

def add_1_up_to_source(left, right, unique_updated_nodes, additions):
    while left < right:
        if left & 1: # if odd, then this is the right child of parent
            additions[left] += 1
            unique_updated_nodes[left] = ref
            left += 1 # iterate AFTER operation
        if right & 1: # if odd, then this is the right child of parent
            right -= 1 # decriment BEFORE operation
            additions[right] += 1
            unique_updated_nodes[right] = ref

        # get to parent by floor division by 2
        left >>= 1
        right >>= 1

def bubble_node_updates_to_maxes(unique_updated_nodes, additions, maxes):
    """
    For each unique updated node, bubble up the max as far as possible (stopping at root)
    """
    for node in unique_updated_nodes:
        node >>= 1
        while node:
            left_child_index = node << 1
            right_child_index = left_child_index + 1

            left_child_max = additions[left_child_index] + maxes[left_child_index]
            right_child_max = additions[right_child_index] + maxes[right_child_index]

            new_max = left_child_max if left_child_max > right_child_max else right_child_max
            # stop early if parent won't get a larger max
            if new_max <= maxes[node]:
                break
            maxes[node] = new_max
            node >>= 1

def update_best_max_at_zoo_index(node, maxes):
    node >>= 1
    # VERY similar to bubble_node_updates_to_maxes
    while node:
        left_child_index = node << 1
        right_child_index = left_child_index + 1

        # ONLY the maxes matter for this update (we're not incorporating any +1 additions)
        left_child_max = maxes[left_child_index]
        right_child_max = maxes[right_child_index]

        new_max = left_child_max if left_child_max > right_child_max else right_child_max
        # stop early if parent won't get a larger max
        if new_max <= maxes[node]:
            break
        maxes[node] = new_max
        node >>= 1

# Reusable components used by all tests
# Must be cleared in between tests 
sources_by_dest_zoo = [[] for _ in ran(10**5)]
sources_by_dest_zoo[0] = ref
cat_unique_updated_nodes = {} # dict is faster than set!
dog_unique_updated_nodes = {}

num_of_case = _int(inp())
for case in ran(num_of_case):
    m_zoos_n_animals_list = inp().split()
    m_zoos = _int(m_zoos_n_animals_list[0])
    n_animals = _int(m_zoos_n_animals_list[1])
    animal_types = inp().split()
    source_zoos = [_int(zoo) for zoo in inp().split()]
    dest_zoos = [_int(zoo) for zoo in inp().split()]

    # Using 1-based indexing, so adding 1 will ensure zoo #1 
    # corresponds to index 1
    zoo_count = m_zoos + 1
    for i in ran(n_animals):
        source_zoo = source_zoos[i]
        dest_zoo = dest_zoos[i]
        # only processing animals with dest_zoo coming AFTER source_zoo 
        if dest_zoo > source_zoo:
            # effectively using counting sort
            sources_by_dest_zoo[dest_zoo].append(source_zoo if animal_types[i] in 'CE' else -source_zoo)
    
    # Establishing segmentation trees
    first_zoo_node = zoo_count + 1
    cat_additions = [0] * (zoo_count << 1)
    cat_maxes     = [*cat_additions] # quicker to just copy
    dog_additions = [*cat_additions]
    dog_maxes     = [*cat_additions]
    # zoo_max_count_list will store the best max values per zoo
    # and later be used to deduce the answer_list

    def sub(sources):
        ### Handle +1 additions at sources ###
        for source in sources:
            if source > 0: # handle cat
                add_1_up_to_source(
                    left = first_zoo_node,
                    right = first_zoo_node + source,
                    unique_updated_nodes = cat_unique_updated_nodes,
                    additions = cat_additions
                )
            else: # handle dog
                add_1_up_to_source(
                    left = first_zoo_node,
                    right = first_zoo_node + (-source), # revert back to pos num
                    unique_updated_nodes = dog_unique_updated_nodes,
                    additions = dog_additions
                )

        # handle updated nodes for both animal types
        bubble_node_updates_to_maxes(
            unique_updated_nodes = cat_unique_updated_nodes,
            additions = cat_additions,
            maxes = cat_maxes
        )
        bubble_node_updates_to_maxes(
            unique_updated_nodes = dog_unique_updated_nodes,
            additions = dog_additions,
            maxes = dog_maxes
        )
        # clear reusable components before next test
        sources.clear() # object reference within sources_by_dest_zoo
        cat_unique_updated_nodes.clear()
        dog_unique_updated_nodes.clear()

        ### Get current max values from both trees ###
        current_cat_max = cat_maxes[1]
        current_dog_max = dog_maxes[1]

        node = zoo_count + zoo_index
        if current_cat_max == current_dog_max:
            # arbitrarily picking cat tree to store zoo_max_list
            cat_maxes[node] = current_cat_max 
            return

        # if the cat tree had the best_max, then update the dog tree at zoo_index
        if current_cat_max > current_dog_max:
            dog_maxes[node] = current_cat_max
            cat_maxes[node] = current_cat_max
            update_best_max_at_zoo_index(node, dog_maxes)
        # else the dog tree had the best_max, then update the cat tree at zoo_index
        else:
            cat_maxes[node] = current_dog_max
            update_best_max_at_zoo_index(node, cat_maxes)

    # MAIN LOOP #
    for zoo_index in ran(1, zoo_count):
        sources = sources_by_dest_zoo[zoo_index]
        # skipping when nothing worth processing
        if sources:
            sub(sources)

    answer_list = [-1] * n_animals
    animal_count = 0
    for zoo_index in ran(1, m_zoos + 1):
        node = zoo_count + zoo_index
        while animal_count < cat_maxes[node]:
            answer_list[animal_count] = zoo_index
            animal_count += 1
    print(*answer_list)