How can I make my solution for Chef and Digits run faster?

Question

Problem Statement

Given \$n\$ digits and \$m\$ indices \$x\$ from \$1. \ldots,n\$, calculate the difference \$b_y = a_x - a_y\$ for all indices \$y \; (y < x)\$. Then, calculate \$B1\$, the sum of all \$b_y\$ which are greater than \$0\$ and \$B2\$, the sum of all \$b_y\$ which are less than \$0\$. The answer for this step is \$B1 - B2\$.

Input

The first line contains two integers \$n, m\$ denoting the number of digits and number of steps. The second line contains \$n\$ digits (without spaces) \$a_1, a_2, ..., a_n\$. Each of next \$m\$ lines contains single integer \$x\$ denoting the index for current step. _{(ADIGIT: Chef and Digits from CodeChef)}

My Issue

I've tried quite hard to optimize the following solution, but online judge always shows Time Limit Exceeded.

My Question

Is it possible to optimize this code further to bring time down less than 1 second, or should I change my approach or perhaps shift to C?

import sys
final=[]
z,y=[int(x) for x in sys.stdin.readline().split()]
stra=sys.stdin.readline()
for i in xrange(y):
    b1=b2=0
    curr=int(input())-1
    for x in stra[:curr]:
        mu=int(stra[curr])-int(x)
        if(mu>0):
            b1+=mu
        else:
            b2+=mu
    final.append(str(b1-b2))
print("\n".join(final))

Answer 1

Issues with your code

You didn't put any effort into choosing good names. Given a code challenge, either stick to the variable names used in the problem statement (n, m, ...) or, preferably, choose explicit and descriptive names such as steps, digits and indices rather than z, y, etc.

Why should you bother with this? Because good names reveal intent, and knowing what the code is supposed to do makes it possible for you and others to understand and improve it easily.

---

Your algorithm uses a naive, brute-force approach. This is fine as a starting point!
But when you get a time limit exceeded message, you need to think about smarter solutions.

• Rene's answer gives you some good advice: He suggests that you employ caching, separate I/O from the calculation, and add up all the absolute values of by, thus skipping the calculation of B1 and B2. This is a good start but you can go much further.
• Janne's answer tells you the essence of what's wrong with your algorithm: you need to avoid looping over the previous digits for every step. Why? Because there will be up to 10⁵ (one hundred thousand) digits and up to 10⁵ steps. However, it is unnecessary to make a sorted copy of the indices — see below.

How?

Start by separating your concerns. Python has functions, so use them to group relevant parts of your code together, while keeping different levels of abstraction separate. Start at the most general level:

def main():
    print(*solve(sys.stdin), sep='\n')

We are saying: print each line of the result of solving the problem with the input from sys.stdin.
Our next concern is to parse the input:

def solve(lines):
    _, digits, *indices = lines
    return calculate_answer(as_ints(digits.rstrip()), as_ints(indices))

We discard the first line _ because we can deduce the length of the digits and the number of steps from the rest of the input, so there is no need to waste cycles parsing them. We then store the digits which are given in the second line, as well as the indices from the following lines.

To get rid of the redundancies in your integer-conversion code, we extract a function:

def as_ints(iterable):
    return [int(x) for x in iterable]

We can pass the digits (stripping the newline character at the end) and the indices to this function to obtain usable values which we can use to calculate_answer(digits, indices). Since we have already parsed the input, this function can be focused on our algorithm to solve the problem.

def calculate_answer(digits, indices):

1. Store the indices in a dictionary, which allows them to be looked up in constant time.

   result = dict.fromkeys(indices)
   

2. Set up another dictionary to count the number of occurrences of each unique digit, which will help us avoid iterating over all previous digits on every step.

   counts = defaultdict(int)
   

3. Iterate over the digits and keep track of each digit's index, starting with 1

   for index, digit in enumerate(digits, start=1):
   

4. If the current index is in result and we have not calculated a value for it yet:
   Assign to the current index the sum of all the absolute values of the differences between the current digit and each unique previous digit d multiplied by its occurrences n.

       if index in result and result[index] is None:
           result[index] = sum(n * abs(digit - d) for d, n in counts.items())
   

5. Increment the number of occurrences of the digit we just encountered:

       counts[digit] += 1
   

6. In the order of the given indices, retrieve each calculated value from result.

   return (result[i] for i in indices)
   

Result

The end result is fast enough to be accepted by CodeChef. As you didn't specify a language version, this is in Python 3, which you should use as well. If you insist on sticking with Python 2.x, you will need to make the necessary adjustments.

import sys
from collections import defaultdict


def main():
    print(*solve(sys.stdin), sep='\n')


def solve(lines):
    _, digits, *indices = lines
    return calculate_answer(as_ints(digits.rstrip()), as_ints(indices))


def calculate_answer(digits, indices):
    result = dict.fromkeys(indices)
    counts = defaultdict(int)
    for index, digit in enumerate(digits, start=1):
        if index in result and result[index] is None:
            result[index] = sum(n * abs(digit - d) for d, n in counts.items())
        counts[digit] += 1
    return (result[i] for i in indices)


def as_ints(iterable):
    return [int(x) for x in iterable]


if __name__ == '__main__':
    main() 

Answer 2

A better algorithm usually wins over micro-optimizations.

• Do you really need to do all these calculations every time? Maybe you could cache them and re-use the results.

• You could perform all the string to int conversions up front, afterwards only looking up numbers from array.

• This whole b1 & b2 thing could be replaced with simple abs() function call. Not sure if it would speed things up, but it sure would simplify the code.

Answer 3

This would avoid looping m times through the string:

1. Sort a copy of the list of indices x given in the input.
2. Sweep through the string, counting the occurrences of each digit as you go along. Each time you hit one of the x, compute the answer from the current counts. Put the answer in a dict indexed by x.
3. Output the answers in correct order by looping through the original list and looking up the value in the dict.