【leetcode 2949 统计美丽子字符串】

import java.util.HashMap;
import java.util.Map;

class Solution {
    public static void main(String[] args) {
        Solution solution = new Solution();
        long ans = solution.beautifulSubstrings("baeyh",2);
        System.out.println(ans);
    }
    public long beautifulSubstrings(String s, int k) {
        int k0 = pSqrt(4 * k);
        Map<String, Integer> map = new HashMap<>();
        long ans = 0;
        int sum = 0;
        map.put((k0-1)+",0",1);
        for (int i = 0; i < s.length(); i++) {
            if (isVowels(s.charAt(i))) {
                sum += 1;
            } else {
                sum -= 1;
            }
            String key = (i%k0) +","+ sum;
            map.put(key,map.getOrDefault(key,0)+1);
            ans += map.get(key) - 1;
        }

        return ans;
    }


    public int pSqrt(int n) {
        int i = 2;
        int res = 1;
        while (i * i <= n) {
            int x = i * i;
            while (n % x == 0) {
                res *= i;
                n = n / x;
            }
            if (n % i == 0) {
                res *= i;
                n = n / i;
            }
            i++;
        }
        if (n != 1) {
            res = res * n;
        }
        return res;
    }

    //
    public boolean isVowels(char ch) {
        if (ch == 'a' || ch == 'e' || ch == 'i' || ch == 'o' || ch == 'u') {
            return true;
        }
        return false;
    }
}