1 / 1
Feb 2018

Okay, this is really frustrating. This is the code I’ve written… but it just gave a TLE. I’m at a loss as to where the inefficiency is! Any ideas anybody? :frowning:

from math import log
from math import ceil

class Seg_Tree:

def __init__(self, leaf_node_count):
    self._build_seg_tree(leaf_node_count)
    self.dirty_list = [0 for _ in range(self.tree_size)]

def print_contents(self):
    for node in self.tree:
        print(node)
    print(self.dirty_list)
    print("====")

def make_node(self, start, end, value):
    node = {}
    node["value"] = value
    node["start"] = start
    node["end"] = end
    return node

def left_index(self, parent_index):
    return parent_index * 2 + 1

def right_index(self, parent_index):
    return parent_index * 2 + 2

def left_node_start(self, parent_index):
    if self.has_left_node(parent_index):
        return self.tree[self.left_index(parent_index)]["start"]
    else:
        return None

def right_node_start(self, parent_index):
    if self.has_right_node(parent_index):
        return self.tree[self.right_index(parent_index)]["start"]
    else:
        return None

def left_node_end(self, parent_index):
    if self.has_left_node(parent_index):
        return self.tree[self.left_index(parent_index)]["end"]
    else:
        return None

def right_node_end(self, parent_index):
    if self.has_right_node(parent_index):
        return self.tree[self.right_index(parent_index)]["end"]
    else:
        return None

def node_start(self, index):
    if self.exists(index):
        return self.tree[index]["start"]
    else:
        return None

def node_end(self, index):
    if self.exists(index):
        return self.tree[index]["end"]
    else:
        return None

def node_value(self, index):
    if self.exists(index):
        return self.tree[index]["value"]
    else:
        return None

def exists(self, index):
    return index < self.tree_size and self.tree[index] != None

def has_left_node(self, index):
    return self.exists(self.left_index(index))

def has_right_node(self, index):
    return self.exists(self.right_index(index))

def no_nodes_match(self, start, end, index):
    return (not self.exists(index)) or end < self.node_start(index) or start > self.node_end(index)

def all_nodes_match(self, start, end, index):
    return self.exists(index) and start <= self.node_start(index) and end >= self.node_end(index)

def some_nodes_match(self, start, end, index):
    return self.exists(index) and (not self.no_nodes_match(start, end, index)) and (not self.all_nodes_match(start, end, index))

def _build_seg_tree(self, leaf_count):
    inner_node_count = 2**ceil(log(leaf_count, 2)) - 1 # This will be a full tree
    self.tree_size = inner_node_count + leaf_count
    self.tree = [None for _ in range(inner_node_count)] #Allocate space for the inner nodes
    for i in range(leaf_count): #Initialise the leaf nodes
        leaf = self.make_node(i, i, 0)
        self.tree.append(leaf)
    for i in range(inner_node_count, -1, - 1): # Now set the values of the inner nodes.
        if self.has_left_node(i):
            start = self.left_node_start(i)
            if self.has_right_node(i):
                end = self.right_node_end(i)
            else:
                end = self.left_node_end(i)
            self.tree[i] = self.make_node(start, end, 0)

def update_node_value(self, index, increment):
    if self.exists(index):
        # Number of nodes in the range X the increment.
        value_to_add = (self.node_end(index) - self.node_start(index) + 1) * increment
        self.tree[index]["value"] += value_to_add
        if self.node_end(index) != self.node_start(index):
            self.dirty_list[index] += increment

def overlap_size(self, start_1, start_2, end_1, end_2):
    start = max(start_1, start_2)
    end = min(end_1, end_2)
    return end - start + 1

def update_range(self, start, end, increment, index=0):
    if self.all_nodes_match(start, end, index):
        self.update_node_value(index, increment)
    elif self.some_nodes_match(start, end, index):
        self.tree[index]["value"] += self.overlap_size(start, self.node_start(index), end, self.node_end(index)) * increment
        self.update_range(start, end, increment, self.left_index(index))
        self.update_range(start, end, increment, self.right_index(index))
    else:
        assert(self.no_nodes_match(start, end, index))

def range_sum(self, start, end, index=0):
    if self.no_nodes_match(start, end, index):
        return 0
    elif self.all_nodes_match(start, end, index):
        return self.node_value(index)
    elif self.some_nodes_match(start, end, index):
        increment = self.dirty_list[index]
        self.dirty_list[index] = 0
        self.update_node_value(self.left_index(index), increment)
        self.update_node_value(self.right_index(index), increment)
        sum_of_left = self.range_sum(start, end, self.left_index(index))
        sum_of_right = self.range_sum(start, end, self.right_index(index))
        return sum_of_left + sum_of_right
    else:
        raise Error("Something went wrong in range sum!")

T = int(input())
for _ in range(T):
N, C = [int(x) for x in input().split()]
tree = Seg_Tree(N)
for c in range©:
command = [int(x) for x in input().split()]
if command[0] == 0:
tree.update_range(command[1] - 1, command[2] - 1, command[3])
else:
print(tree.range_sum(command[1] - 1, command[2] - 1))