Problem Statement:

You are given a function \$f(x)\$, where \$f(x)\$ is \$1\$ if the first and last characters of string \$x\$ are equal; else it is \$0\$. You are given a string \$S\$ and you have to find the sum of \$f(x)\$ for all substrings \$x\$ of given string \$S\$.

Sample Input:


Sample Output:



f("a")=1, f("aba")=1, f("abaca")=1 but f("ab")=0, f("bac")=0. Hence counting all substrings we get 14.

The 14 substring are

a - 4(times) 
b - 2 
c - 1 
aba - 2 
bab - 1 
aca - 1 
ababa - 1 
abaca - 1 
ababaca - 1

My code:

l = int(input())
s = [x for x in input().strip()]
print(l + sum([1 for i in range(0,l-1) for j in range(i+1,l) if s[i] == s[j]]))

I have tried the following snippets too:


from collections import Counter
import math
l, s = int(input()), [x for x in input()]
print(l + sum([math.factorial(value-1) for value in Counter(s).values() if value != 1]))

Wrong answer:

import itertools
l, s = int(input()), input()
print(sum(int(i[0] == i[-1]) for i in itertools.combinations_with_replacement(s,2)))

The solutions is working fine except that it times out in certain test cases. Can anyone suggest a better way to do this?


5 Answers 5


First review

It is almost always a good idea to put your logic in a function (or class) that can be easily documented and tested. That way, in anything goes wrong during the optimisation, you'll know it straight-away. Also, this can be useful if you need to measure your result to ensure your optimisation is an actual optimisation.

Just moving the code around, here is what I start with :

def get_nb_substring(s):
    l = len(s)
    return(l + sum([1 for i in range(0,l-1) for j in range(i+1,l) if s[i] == s[j]]))

def test_from_input():
    _ = int(input()) # useful ?
    s = [x for x in input().strip()]

def automatic_tests():
    assert get_nb_substring("ababaca") == 14

if __name__ == '__main__':

Now, the actual review may start.

  • you don't need 0 as a first argument for range.

  • you usually don't need the length when you are using iterations in Python. When you do need it, it usually means you are doing something wrong.

  • for that reason, what you are doing with indices can be achieved with itertools.combinations(s, 2)

Now, the code can be rewritten :

def get_nb_substring(s):
    return len(s) + sum(1 for i, j in itertools.combinations(s, 2) if i == j)

which is still 0(n^2).


An idea could be to check where the different letters appear. Then if a letter c appears at positions p1, p2, ..., pn, you know it will contribute for substrings : p1, p1-p2, p1-p3, ..., p1-pn, p2,p2-p3, p2-p4, ... p2-pn, .... pn. There will be 1+2+...+n = n*(n+1)/2 such substrings.

There, if I did everything correctly, I end up with O(n) code :

def get_nb_substring(s):
    position = dict()
    for i, c in enumerate(s):
        position.setdefault(c, []).append(i)
    return sum(l*(l+1)/2 for l in (len (l) for l in position.values()))

On the test example, it seems to work. On other cases, I obtain different results than your code so I don't know which is correct.

I've used a dict to map characters to a list of position but a mapping from characters to number of positions would have been enough. This corresponds to using the Counter collection)

def get_nb_substring(s):
    return sum(v*(v+1)/2 for v in collections.Counter(s).values())

If you append some character c to a string, how many new substrings can you make that start with c and end with c? As many as the number of c you have seen so far.

You can iterate over the characters of the input string and build a map of counts. Before incrementing the count for a character, add to the running sum the current count. In case of ababaca, this will happen:

  • a : current count 0, set it to 1
  • b : current count 0, set it to 1
  • a : current count 1, add it to the sum, increment count
  • b : current count 1, add it to the sum, increment count
  • a : current count 2, add it to the sum, increment count
  • c : current count 0, set it to 1
  • a : current count 3, add it to the sum, increment count

The total sum at this point is 7, add to this the number of characters. Time complexity \$O(N)\$

Btw, there's no need to make a list of characters like this:

s = [x for x in input().strip()]

If you do just this your program will work just the same:

s = input().strip()

Putting it together:

from collections import Counter

def count_begin_end(word):
    total = len(word)
    counter = Counter()
    for c in word:
        total += counter[c]
        counter[c] += 1
    return total

def main():
    _, word = input(), input()

if __name__ == '__main__':

You're creating a few lists where you don't need them. I don't know if this fixes the timeout problem, but here's my take:

def countF(s):
    return sum(int(s[i]==s[j-1]) for i in range(l) for j in range(i+1,l+1))

print(countF("ababaca")) #14

First of all, I'm not transforming the string into a list — it's not needed, you can get to a specific characters using square brackets just like you could in a list.

Second, instead of a list comprehension, I'm using a generator comprehension — this way no list is built up in memory. For example, given a very long string, your code needs to build a list the size of the cardinality of possible substrings, which can get quite large. Using my generator, it sums up while looping, so it's always just a single number it needs to keep in memory.

edit: Good point about the slicing, this way should be better.

  • \$\begingroup\$ Hey.. Appreciate your reply. But can you explain how this solution is different to mine? \$\endgroup\$
    – gaurav
    Commented Jun 20, 2015 at 22:11
  • \$\begingroup\$ Also, wouldn't s[i:j] create a substring which 1) Is an expensive command 2) is unneeded since we only want to compare the first and last character of the substring? \$\endgroup\$
    – gaurav
    Commented Jun 20, 2015 at 22:12
  • \$\begingroup\$ @GauravKeswani Was just editing it. Is this enough? \$\endgroup\$
    – L3viathan
    Commented Jun 20, 2015 at 22:12
  • \$\begingroup\$ The generator comprehension bit makes sense to me. But can you explain why we are creating the substring? It isn't needed right? \$\endgroup\$
    – gaurav
    Commented Jun 20, 2015 at 22:13
  • \$\begingroup\$ Good point, my edit should make it faster. int isn't really needed, but I doubt it makes it any slower either. \$\endgroup\$
    – L3viathan
    Commented Jun 20, 2015 at 22:14

Since this is something totally different, I have another answer, which should be in \$O(N)\$. With standard lib, but all unoptimized at the moment and just proof of concept:

from collections import Counter
for ch in s:
  • \$\begingroup\$ I was thinking on similar lines and had come up with this some time back from collections import Counter import math l, s = int(input()), [x for x in input()] print(l + sum([math.factorial(value-1) for value in Counter(s).values() if value != 1])) \$\endgroup\$
    – gaurav
    Commented Jun 20, 2015 at 22:36

After numerous suggestions on this portal and a good night's sleep, I could come up with this :)

from collections import Counter
_, s = int(input()), input()
occurences = Counter(s).values()
print(sum(n*(n+1)//2 for n in occurences))

And it works! :D

enter image description here

  • \$\begingroup\$ Is this an O(1) solution or a O(NlogN) solution? I think it is nlogn but I need confirmation. \$\endgroup\$
    – gaurav
    Commented Jun 21, 2015 at 20:13

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.