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


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


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

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

if __name__ == '__main__':
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):
l=len(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.

• Hey.. Appreciate your reply. But can you explain how this solution is different to mine? – Gaurav Keswani Jun 20 '15 at 22:11
• 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? – Gaurav Keswani Jun 20 '15 at 22:12
• @GauravKeswani Was just editing it. Is this enough? – L3viathan Jun 20 '15 at 22:12
• The generator comprehension bit makes sense to me. But can you explain why we are creating the substring? It isn't needed right? – Gaurav Keswani Jun 20 '15 at 22:13
• Good point, my edit should make it faster. int isn't really needed, but I doubt it makes it any slower either. – L3viathan Jun 20 '15 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
s="ababaca"
i=7
c=Counter()
for ch in s:
c[ch]+=1
n+=c[ch]-1
print(n+i)

• 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])) – Gaurav Keswani Jun 20 '15 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

• Is this an O(1) solution or a O(NlogN) solution? I think it is nlogn but I need confirmation. – Gaurav Keswani Jun 21 '15 at 20:13