distinct_integers_in_range.py

# source: https://www.hackerearth.com/practice/data-structures/advanced-data-structures/segment-trees/practice-problems/algorithm/distinct-integers-in-range-66eca44b/description/
import re

### SETUP ###
max_int = 5000
MAX_VALUE = (1 << max_int) - 1
re_splitter = re.compile(" ")

def agg_func(val_1, val_2):
    output = val_1 | val_2
    if output == MAX_VALUE:
        return -1
    return output

### HANDLING SEGMENTATION TREES ###
def init_update(left, right, seg_node_index = 0):
    if left > right:
        return

    if left == right:
        a_seg_tree[seg_node_index] = a_input_list[left]
        b_seg_tree[seg_node_index] = b_input_list[left]
        return
    
    midpoint = (left + right) >> 1
    
    left_seg_node_index = (seg_node_index << 1) + 1
    init_update(left, midpoint, left_seg_node_index)

    right_seg_node_index = (seg_node_index << 1) + 2
    init_update(midpoint + 1, right, right_seg_node_index)

    a_seg_tree[seg_node_index] = agg_func(a_seg_tree[left_seg_node_index], a_seg_tree[right_seg_node_index])
    
    b_seg_tree[seg_node_index] = agg_func(b_seg_tree[left_seg_node_index], b_seg_tree[right_seg_node_index])

def a_query(query_left_index, query_right_index, left, right, seg_node_index = 0):
    if left > right or query_right_index < left or right < query_left_index:
        return 0
    if query_left_index <= left and right <= query_right_index:
        return a_seg_tree[seg_node_index]
    
    midpoint = (left + right) >> 1
    
    left_value = a_query(query_left_index, query_right_index, left, midpoint, (seg_node_index << 1) + 1)
    if left_value < 0:
        return -1

    right_value = a_query(query_left_index, query_right_index, midpoint + 1, right, (seg_node_index << 1) + 2)

    return agg_func(left_value, right_value)

def b_query(query_left_index, query_right_index, left, right, seg_node_index = 0):
    if left > right or query_right_index < left or right < query_left_index:
        return 0
    if query_left_index <= left and right <= query_right_index:
        return b_seg_tree[seg_node_index]
    
    midpoint = (left + right) >> 1
    
    left_value = b_query(query_left_index, query_right_index, left, midpoint, (seg_node_index << 1) + 1)
    if left_value < 0:
        return -1

    right_value = b_query(query_left_index, query_right_index, midpoint + 1, right, (seg_node_index << 1) + 2)

    return agg_func(left_value, right_value)

def convert_input(string_num):
    return 1 << (int(string_num) - 1)

n_leaf_nodes = int(input())

a_input_list = tuple(map(convert_input, re_splitter.split(input())))
b_input_list = tuple(map(convert_input, re_splitter.split(input())))

max_size = 262143 # hard coding for efficiency

a_seg_tree = [0] * max_size
b_seg_tree = [0] * max_size

init_update(0, n_leaf_nodes - 1)

### HANDLING QUERIES ###
count_memo = {-1: max_int}
    
for _ in range(int(input())):
    a_start, a_end, b_start, b_end = map(int, re_splitter.split(input()))

    set_a = a_query(a_start - 1, a_end - 1, 0, n_leaf_nodes - 1)
    if set_a < 0:
        print(max_int)
        continue
    
    set_b = b_query(b_start - 1, b_end - 1, 0, n_leaf_nodes - 1)
    if set_b < 0:
        print(max_int)
        continue

    result = agg_func(set_a, set_b)
    count = count_memo.get(result)
    if count:
        print(count)
    else:
        # most efficient way of counting bits
        count = bin(result).count("1")
        count_memo[result] = count
        print(count)