<p>You are given a <strong>0-indexed </strong>integer array <code>nums</code>.</p>

<p>The <strong>distinct count</strong> of a subarray of <code>nums</code> is defined as:</p>

<ul>
	<li>Let <code>nums[i..j]</code> be a subarray of <code>nums</code> consisting of all the indices from <code>i</code> to <code>j</code> such that <code>0 &lt;= i &lt;= j &lt; nums.length</code>. Then the number of distinct values in <code>nums[i..j]</code> is called the distinct count of <code>nums[i..j]</code>.</li>
</ul>

<p>Return <em>the sum of the <strong>squares</strong> of <strong>distinct counts</strong> of all subarrays of </em><code>nums</code>.</p>

<p>A subarray is a contiguous <strong>non-empty</strong> sequence of elements within an array.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,2,1]
<strong>Output:</strong> 15
<strong>Explanation:</strong> Six possible subarrays are:
[1]: 1 distinct value
[2]: 1 distinct value
[1]: 1 distinct value
[1,2]: 2 distinct values
[2,1]: 2 distinct values
[1,2,1]: 2 distinct values
The sum of the squares of the distinct counts in all subarrays is equal to 1<sup>2</sup> + 1<sup>2</sup> + 1<sup>2</sup> + 2<sup>2</sup> + 2<sup>2</sup> + 2<sup>2</sup> = 15.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,1]
<strong>Output:</strong> 3
<strong>Explanation:</strong> Three possible subarrays are:
[1]: 1 distinct value
[1]: 1 distinct value
[1,1]: 1 distinct value
The sum of the squares of the distinct counts in all subarrays is equal to 1<sup>2</sup> + 1<sup>2</sup> + 1<sup>2</sup> = 3.</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= nums.length &lt;= 100</code></li>
	<li><code>1 &lt;= nums[i] &lt;= 100</code></li>
</ul>

from collections import Counter

class Solution:
    def __init__(self):
        self._nums: List[int]
        self._length: int

    def setNums(self, nums: List[int]) -> None:
        """Sets nums and its length."""
        self._nums = nums

    def sumCounts(self, nums: List[int]) -> int:
        self.setNums(nums)
        self._length = len(nums)

        counts = 0

        for window_size in range(1, self._length + 1):
            print(f'Window size: {window_size}')
            counts += self._findDistinct(window_size)
        return counts

    def _findDistinct(self, window_size: int) -> int:
        """Finds the distinct count of all subarrays of a window_size."""
        start_index = self._length - window_size
        end_index = start_index + window_size - 1
        # print(f'\tWindow starting at {start_index} and ending at {end_index}.')
        subarray = self._nums[ start_index: self._length].copy()
        subarray = Counter(subarray)

        iterations = self._length + 1 - window_size  # i.e. If window_size = 1 >>> Iterations == length (1 per)
        counts = 0

        distinct = len(subarray.keys())
        counts += distinct ** 2
        # print(f'\t COUNTED {subarray.keys()}')
        for i in range(iterations - 1):
            start_index -= 1
            add_value = self._nums[start_index]
            subarray[add_value] += 1

            remove_value = self._nums[end_index]
            subarray[remove_value] -= 1
            if subarray[remove_value] == 0:
                del subarray[remove_value]
            end_index -= 1

            # print(f'\t\t\t Value to add: {add_value}; value to remove {remove_value}')
            distinct = len(subarray.keys())
            counts += distinct ** 2
            # print(f'\t COUNTED {subarray.keys()}')
        return counts



198