Number of sub strings with same first and last character 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:

7
ababaca


Sample Output:

14


Explanation: 

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:

Timeout:

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?
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()]
    print(get_nb_substring(s))

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

if __name__ == '__main__':
    automatic_tests()


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

Optimisations

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