Number of sub strings with same first and last character

Question

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?

Answer 1

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

Answer 2

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()
    print(count_begin_end(word))


if __name__ == '__main__':
    main()

Answer 3

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.

Answer 4

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)

Answer 5

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