aboutsummaryrefslogtreecommitdiff
path: root/subarr_distinct_elem_sum.c
blob: 6c01ad81298f80796d1611c769f97e430ba5c8bd (plain)
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
}