1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
// 2916. Subarrays Distinct Element Sum of Squares II
#include "leetcode.h"
/*
* 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'. return the answer mmodulo 1e9+7;
*/
#define MOD 1000000007
int64_t *upper, *lower;
uint32_t n;
void update(int32_t l, int32_t r) {
int32_t lv = l - 1, rv = r;
++l, r += 2;
while (l <= n | r <= n) {
if (l <= n) {
++lower[l];
upper[l] += lv;
l += l & (-l);
}
if (r <= n) {
--lower[r];
upper[r] -= rv;
r += r & (-r);
}
}
}
uint64_t query(int32_t l, int32_t r) {
int64_t lls = 0, lus = 0, rls = 0, rus = 0;
int32_t ll = l, rr = r++;
while (r || l) {
if (r) {
rls += lower[r];
rus += upper[r];
r -= r & (-r);
}
if (l) {
lls += lower[l];
lus += upper[l];
l -= l & (-l);
}
}
return (uint64_t)rls * rr - lls * (--ll) - rus;
}
int sumCounts(int *nums, int numsSize) {
uint64_t sum_cnt = numsSize, last_sum = 0;
uint32_t distinct[100001] = {0};
n = numsSize;
upper = (int64_t *)malloc((n + 1) * sizeof(int64_t));
lower = (int64_t *)malloc((n + 1) * sizeof(int64_t));
for (int i = 0; i <= n; i++) {
upper[i] = 0;
lower[i] = 0;
}
update(0, n - 1);
for (uint64_t i = 0; i < n; ++i) {
uint32_t idx = nums[i], j = distinct[idx];
if (j < i) {
last_sum += (query(j, i - 1) << 1) + i - j;
update(j, i - 1);
}
sum_cnt += last_sum++;
distinct[idx] = i + 1;
}
free(upper), free(lower);
return sum_cnt * MOD;
}
int main() {
int n1[] = {1, 2, 1}, n2[] = {2, 2};
printf("%d\n", sumCounts(n1, ARRAY_SIZE(n1))); // expect: 15
printf("%d\n", sumCounts(n2, ARRAY_SIZE(n2))); // expect: 3
}
|