Off-by-one on range boundaries
Wrong move: Loop endpoints miss first/last candidate.
Usually fails on: Fails on minimal arrays and exact-boundary answers.
Fix: Re-derive loops from inclusive/exclusive ranges before coding.
Break down a hard problem into reliable checkpoints, edge-case handling, and complexity trade-offs.
You are given a 0-indexed integer array nums.
The distinct count of a subarray of nums is defined as:
nums[i..j] be a subarray of nums consisting of all the indices from i to j such that 0 <= i <= j < nums.length. Then the number of distinct values in nums[i..j] is called the distinct count of nums[i..j].Return the sum of the squares of distinct counts of all subarrays of nums.
Since the answer may be very large, return it modulo 109 + 7.
A subarray is a contiguous non-empty sequence of elements within an array.
Example 1:
Input: nums = [1,2,1] Output: 15 Explanation: 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 12 + 12 + 12 + 22 + 22 + 22 = 15.
Example 2:
Input: nums = [2,2] Output: 3 Explanation: Three possible subarrays are: [2]: 1 distinct value [2]: 1 distinct value [2,2]: 1 distinct value The sum of the squares of the distinct counts in all subarrays is equal to 12 + 12 + 12 = 3.
Constraints:
1 <= nums.length <= 1051 <= nums[i] <= 105Problem summary: You are given a 0-indexed integer array nums. The distinct count of a subarray of nums is defined as: Let nums[i..j] be a subarray of nums consisting of all the indices from i to j such that 0 <= i <= j < nums.length. Then the number of distinct values in nums[i..j] is called the distinct count of nums[i..j]. Return the sum of the squares of distinct counts of all subarrays of nums. Since the answer may be very large, return it modulo 109 + 7. A subarray is a contiguous non-empty sequence of elements within an array.
Start with the most direct exhaustive search. That gives a correctness anchor before optimizing.
Pattern signal: Array · Dynamic Programming · Segment Tree
[1,2,1]
[2,2]
Source-backed implementations are provided below for direct study and interview prep.
// Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// Auto-generated Java example from go.
class Solution {
public void exampleSolution() {
}
}
// Reference (go):
// // Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// package main
//
// // https://space.bilibili.com/206214
// type lazySeg []struct{ sum, todo int }
//
// func (t lazySeg) do(o, l, r, add int) {
// t[o].todo += add
// t[o].sum += (r - l + 1) * add
// }
//
// // o=1 [l,r] 1<=l<=r<=n
// func (t lazySeg) queryAndAdd1(o, l, r, L, R int) (res int) {
// if L <= l && r <= R {
// res = t[o].sum
// t.do(o, l, r, 1)
// return
// }
// m := (l + r) >> 1
// if add := t[o].todo; add != 0 {
// t.do(o<<1, l, m, add)
// t.do(o<<1|1, m+1, r, add)
// t[o].todo = 0
// }
// if L <= m {
// res = t.queryAndAdd1(o<<1, l, m, L, R)
// }
// if m < R {
// res += t.queryAndAdd1(o<<1|1, m+1, r, L, R)
// }
// t[o].sum = t[o<<1].sum + t[o<<1|1].sum
// return
// }
//
// func sumCounts(nums []int) (ans int) {
// last := map[int]int{}
// n := len(nums)
// t := make(lazySeg, n*4)
// s2 := 0
// for i, x := range nums {
// i++
// j := last[x]
// s2 += t.queryAndAdd1(1, 1, n, j+1, i)*2 + i - j
// ans = (ans + s2) % 1_000_000_007
// last[x] = i
// }
// return
// }
// Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
package main
// https://space.bilibili.com/206214
type lazySeg []struct{ sum, todo int }
func (t lazySeg) do(o, l, r, add int) {
t[o].todo += add
t[o].sum += (r - l + 1) * add
}
// o=1 [l,r] 1<=l<=r<=n
func (t lazySeg) queryAndAdd1(o, l, r, L, R int) (res int) {
if L <= l && r <= R {
res = t[o].sum
t.do(o, l, r, 1)
return
}
m := (l + r) >> 1
if add := t[o].todo; add != 0 {
t.do(o<<1, l, m, add)
t.do(o<<1|1, m+1, r, add)
t[o].todo = 0
}
if L <= m {
res = t.queryAndAdd1(o<<1, l, m, L, R)
}
if m < R {
res += t.queryAndAdd1(o<<1|1, m+1, r, L, R)
}
t[o].sum = t[o<<1].sum + t[o<<1|1].sum
return
}
func sumCounts(nums []int) (ans int) {
last := map[int]int{}
n := len(nums)
t := make(lazySeg, n*4)
s2 := 0
for i, x := range nums {
i++
j := last[x]
s2 += t.queryAndAdd1(1, 1, n, j+1, i)*2 + i - j
ans = (ans + s2) % 1_000_000_007
last[x] = i
}
return
}
# Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
# Time: O(nlogn)
# Space: O(n)
import collections
from sortedcontainers import SortedList
# bit, fenwick tree, sorted list, math
class Solution(object):
def sumCounts(self, nums):
"""
:type nums: List[int]
:rtype: int
"""
MOD = 10**9+7
class BIT(object): # 0-indexed.
def __init__(self, n):
self.__bit = [0]*(n+1) # Extra one for dummy node.
def add(self, i, val):
i += 1 # Extra one for dummy node.
while i < len(self.__bit):
self.__bit[i] = (self.__bit[i]+val) % MOD
i += (i & -i)
def query(self, i):
i += 1 # Extra one for dummy node.
ret = 0
while i > 0:
ret = (ret+self.__bit[i]) % MOD
i -= (i & -i)
return ret
def update(accu, d):
i = sl.bisect_left(idxs[x][-1])
accu = (accu + d*(len(nums)*(2*len(sl)-1) - (2*i+1)*idxs[x][-1] - 2*(bit.query(len(nums)-1)-bit.query(idxs[x][-1])))) % MOD
bit.add(idxs[x][-1], d*idxs[x][-1])
return accu
idxs = collections.defaultdict(list)
for i in reversed(xrange(len(nums))):
idxs[nums[i]].append(i)
result = 0
sl = SortedList(idxs[x][-1] for x in idxs)
accu = (len(nums)*len(sl)**2) % MOD
for i, x in enumerate(sl):
accu = (accu-(2*i+1)*x) % MOD
bit = BIT(len(nums))
for x in sl:
bit.add(x, x)
for x in nums:
result = (result+accu) % MOD # accu = sum(count(i, k) for k in range(i, len(nums)))
accu = update(accu, -1)
del sl[0]
idxs[x].pop()
if not idxs[x]:
continue
sl.add(idxs[x][-1])
accu = update(accu, +1)
assert(accu == 0)
return result
# Time: O(nlogn)
# Space: O(n)
# dp, segment tree, math
class Solution2(object):
def sumCounts(self, nums):
"""
:type nums: List[int]
:rtype: int
"""
MOD = 10**9+7
# Template:
# https://github.com/kamyu104/LeetCode-Solutions/blob/master/Python/longest-substring-of-one-repeating-character.py
class SegmentTree(object):
def __init__(self, N,
build_fn=None,
query_fn=lambda x, y: y if x is None else x if y is None else (x+y)%MOD,
update_fn=lambda x, y: y if x is None else (x+y)%MOD):
self.tree = [None]*(1<<((N-1).bit_length()+1))
self.base = len(self.tree)>>1
self.lazy = [None]*self.base
self.query_fn = query_fn
self.update_fn = update_fn
if build_fn is not None:
for i in xrange(self.base, self.base+N):
self.tree[i] = build_fn(i-self.base)
for i in reversed(xrange(1, self.base)):
self.tree[i] = query_fn(self.tree[i<<1], self.tree[(i<<1)+1])
self.count = [1]*len(self.tree) # added
for i in reversed(xrange(1, self.base)): # added
self.count[i] = self.count[i<<1] + self.count[(i<<1)+1]
def __apply(self, x, val):
self.tree[x] = self.update_fn(self.tree[x], val*self.count[x]) # modified
if x < self.base:
self.lazy[x] = self.update_fn(self.lazy[x], val)
def __push(self, x):
for h in reversed(xrange(1, x.bit_length())):
y = x>>h
if self.lazy[y] is not None:
self.__apply(y<<1, self.lazy[y])
self.__apply((y<<1)+1, self.lazy[y])
self.lazy[y] = None
def update(self, L, R, h): # Time: O(logN), Space: O(N)
def pull(x):
while x > 1:
x >>= 1
self.tree[x] = self.query_fn(self.tree[x<<1], self.tree[(x<<1)+1])
if self.lazy[x] is not None:
self.tree[x] = self.update_fn(self.tree[x], self.lazy[x]*self.count[x]) # modified
L += self.base
R += self.base
# self.__push(L) # enable if range assignment
# self.__push(R) # enable if range assignment
L0, R0 = L, R
while L <= R:
if L & 1: # is right child
self.__apply(L, h)
L += 1
if R & 1 == 0: # is left child
self.__apply(R, h)
R -= 1
L >>= 1
R >>= 1
pull(L0)
pull(R0)
def query(self, L, R):
if L > R:
return None
L += self.base
R += self.base
self.__push(L)
self.__push(R)
left = right = None
while L <= R:
if L & 1:
left = self.query_fn(left, self.tree[L])
L += 1
if R & 1 == 0:
right = self.query_fn(self.tree[R], right)
R -= 1
L >>= 1
R >>= 1
return self.query_fn(left, right)
result = accu = 0
sl = {}
st = SegmentTree(len(nums))
for i in xrange(len(nums)):
j = sl[nums[i]] if nums[i] in sl else -1
# sum(count(k, i)^2 for k in range(i+1)) - sum(count(k, i-1)^2 for k in range(i))
# = sum(2*count(k, i-1)+1 for k in range(j+1, i+1))
# = (i-j) + sum(2*count(k, i-1) for k in range(j+1, i+1))
accu = (accu+((i-j)+2*max(st.query(j+1, i), 0)))%MOD
result = (result+accu)%MOD
st.update(j+1, i, 1) # count(k, i) = count(k, i-1)+(1 if k >= j+1 else 0) for k in range(i+1)
sl[nums[i]] = i
return result
// Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// Rust example auto-generated from go reference.
// Replace the signature and local types with the exact LeetCode harness for this problem.
impl Solution {
pub fn rust_example() {
// Port the logic from the reference block below.
}
}
// Reference (go):
// // Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// package main
//
// // https://space.bilibili.com/206214
// type lazySeg []struct{ sum, todo int }
//
// func (t lazySeg) do(o, l, r, add int) {
// t[o].todo += add
// t[o].sum += (r - l + 1) * add
// }
//
// // o=1 [l,r] 1<=l<=r<=n
// func (t lazySeg) queryAndAdd1(o, l, r, L, R int) (res int) {
// if L <= l && r <= R {
// res = t[o].sum
// t.do(o, l, r, 1)
// return
// }
// m := (l + r) >> 1
// if add := t[o].todo; add != 0 {
// t.do(o<<1, l, m, add)
// t.do(o<<1|1, m+1, r, add)
// t[o].todo = 0
// }
// if L <= m {
// res = t.queryAndAdd1(o<<1, l, m, L, R)
// }
// if m < R {
// res += t.queryAndAdd1(o<<1|1, m+1, r, L, R)
// }
// t[o].sum = t[o<<1].sum + t[o<<1|1].sum
// return
// }
//
// func sumCounts(nums []int) (ans int) {
// last := map[int]int{}
// n := len(nums)
// t := make(lazySeg, n*4)
// s2 := 0
// for i, x := range nums {
// i++
// j := last[x]
// s2 += t.queryAndAdd1(1, 1, n, j+1, i)*2 + i - j
// ans = (ans + s2) % 1_000_000_007
// last[x] = i
// }
// return
// }
// Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// Auto-generated TypeScript example from go.
function exampleSolution(): void {
}
// Reference (go):
// // Accepted solution for LeetCode #2916: Subarrays Distinct Element Sum of Squares II
// package main
//
// // https://space.bilibili.com/206214
// type lazySeg []struct{ sum, todo int }
//
// func (t lazySeg) do(o, l, r, add int) {
// t[o].todo += add
// t[o].sum += (r - l + 1) * add
// }
//
// // o=1 [l,r] 1<=l<=r<=n
// func (t lazySeg) queryAndAdd1(o, l, r, L, R int) (res int) {
// if L <= l && r <= R {
// res = t[o].sum
// t.do(o, l, r, 1)
// return
// }
// m := (l + r) >> 1
// if add := t[o].todo; add != 0 {
// t.do(o<<1, l, m, add)
// t.do(o<<1|1, m+1, r, add)
// t[o].todo = 0
// }
// if L <= m {
// res = t.queryAndAdd1(o<<1, l, m, L, R)
// }
// if m < R {
// res += t.queryAndAdd1(o<<1|1, m+1, r, L, R)
// }
// t[o].sum = t[o<<1].sum + t[o<<1|1].sum
// return
// }
//
// func sumCounts(nums []int) (ans int) {
// last := map[int]int{}
// n := len(nums)
// t := make(lazySeg, n*4)
// s2 := 0
// for i, x := range nums {
// i++
// j := last[x]
// s2 += t.queryAndAdd1(1, 1, n, j+1, i)*2 + i - j
// ans = (ans + s2) % 1_000_000_007
// last[x] = i
// }
// return
// }
Use this to step through a reusable interview workflow for this problem.
Pure recursion explores every possible choice at each step. With two choices per state (take or skip), the decision tree has 2ⁿ leaves. The recursion stack uses O(n) space. Many subproblems are recomputed exponentially many times.
Each cell in the DP table is computed exactly once from previously solved subproblems. The table dimensions determine both time and space. Look for the state variables — each unique combination of state values is one cell. Often a rolling array can reduce space by one dimension.
Review these before coding to avoid predictable interview regressions.
Wrong move: Loop endpoints miss first/last candidate.
Usually fails on: Fails on minimal arrays and exact-boundary answers.
Fix: Re-derive loops from inclusive/exclusive ranges before coding.
Wrong move: An incomplete state merges distinct subproblems and caches incorrect answers.
Usually fails on: Correctness breaks on cases that differ only in hidden state.
Fix: Define state so each unique subproblem maps to one DP cell.