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

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 (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))

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 105 (one hundred thousand) digits and up to 105 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):

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

result = dict.fromkeys(indices)

-
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)

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

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

-
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())

-
Increment the number of occurrences of the digit we just encountered:

counts[digit] += 1

-
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()

StackExchange Code Review Q#46320, answer score: 10