LeetCode in Kotlin

2416. Sum of Prefix Scores of Strings

Hard

You are given an array words of size n consisting of non-empty strings.

We define the score of a string word as the number of strings words[i] such that word is a prefix of words[i].

Return an array answer of size n where answer[i] is the sum of scores of every non-empty prefix of words[i].

Note that a string is considered as a prefix of itself.

Example 1:

Input: words = [“abc”,”ab”,”bc”,”b”]

Output: [5,4,3,2]

Explanation: The answer for each string is the following:

The total is answer[0] = 2 + 2 + 1 = 5.

The total is answer[1] = 2 + 2 = 4.

The total is answer[2] = 2 + 1 = 3.

The total is answer[3] = 2.

Example 2:

Input: words = [“abcd”]

Output: [4]

Explanation:

“abcd” has 4 prefixes: “a”, “ab”, “abc”, and “abcd”.

Each prefix has a score of one, so the total is answer[0] = 1 + 1 + 1 + 1 = 4.

Constraints:

Solution

class Solution {
    private val child: Array<Solution?> = arrayOfNulls(26)
    private var ct: Int = 0

    fun sumPrefixScores(words: Array<String>): IntArray {
        for (s in words) {
            insert(s)
        }
        val res = IntArray(words.size)
        for (i in words.indices) {
            val word = words[i]
            res[i] = countPre(word)
        }
        return res
    }

    private fun insert(word: String) {
        var cur: Solution? = this
        for (element in word) {
            val id = element.code - 'a'.code
            if (cur!!.child[id] == null) {
                cur.child[id] = Solution()
            }
            cur.child[id]!!.ct++
            cur = cur.child[id]
        }
    }

    private fun countPre(word: String): Int {
        var cur: Solution? = this
        var localCt = 0
        for (element in word) {
            val id = element.code - 'a'.code
            if (cur!!.child[id] == null) {
                return localCt
            }
            localCt += cur.ct
            cur = cur.child[id]
        }
        localCt += cur!!.ct
        return localCt
    }
}