Number of substrings starting with a vowel

number_of_substrings_starting_with_a_vowel.py

class Solution(object):

    def solve(self, s):

        vowels = {"a", "e", "i", "o", "u", "A", "E", "I", "O", "U"}

        ans = 0

        for i, c in enumerate(s):
            if c in vowels:
               ans += len(s) - i

        return ans

if __name__ == "__main__":
    a = "ABEC"
    ans = Solution().solve(a)
    print(ans)