Optimise array calculation in python

64 Views Asked by At

The task:

You are given an array a of 0 < N ≤ 100000 elements and 0 < M ≤ 100000 queries. Each query can be one of two types:

? l r k — REQ(l, r, k)

or

+ i delta — inc(i, delta): increase a[i] value by delta.

It is guaranteed that elements of array never exceed K, that is, 0 ≤ a[i] ≤ K.

Input format:

  • The first line contains 3 integers N, M, and K, where 0 < N, M ≤ 100000 and 0 < K ≤ 100. N is the number of items in the array, M is the number of queries and K is an upper bound on the values in the array.

  • The next line contains N integer numbers — the values of the array a, where 0 ≤ a[i] ≤ K

  • Each of next M lines contains one character followed by 3 or 4 integer numbers — the request can be one of two types:

? l r k — REQ(l, r, k): 0 ≤ l < r ≤ N, 0 ≤ k ≤ K

or

+ i delta — inc(i, delta): increase a[i] value by delta, 0 ≤ i < N, -K ≤ k ≤ K

Output format:

For each "?" query provide answer for REQ request.

Input sample:

2 8 1
0 0
? 0 2 0
? 0 2 1
+ 0 1
? 0 2 0
? 0 2 1
+ 1 1
? 0 2 0
? 0 2 1

Output:

2
0
1
1
0
2

The code:

import sys


class FenwickTree:
    def __init__(self, n, k, tree):
        self.n = n
        self.k = k
        self.tree = tree

    def add(self, i, delta):
        self.tree[i] += delta

    def sum(self, fr, to, val):
        return self.tree[fr:to].count(val)


if __name__ == '__main__':
    sys.stdin = open('input.txt', 'r')
    n, m, k = map(int, input().split())
    a = list(map(int, input().split()))

    ft = FenwickTree(n, k, a)

    for _ in range(m):
        op = input().split()
        if op[0] == '?':
            print(ft.sum(int(op[1]), int(op[2]), int(op[3])))
        else:
            ft.add(int(op[1]), int(op[2]))

I look to optimise the code above, as it doesn't fit into a time limit.

0

There are 0 best solutions below