aboutsummaryrefslogtreecommitdiff
path: root/find_sum_arr_product_seq.c
blob: 568eb093648e8f1570ec7b00bc76a6338c34ed96 (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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
// 3539. Find Sum of Array Product of Magical Sequences
#include "leetcode.h"

/*
 * you are given two integers 'm' and 'k', and an integer array 'nums'. a
 * sequence of integers 'seq' is called magical if 'seq' has a size of 'm', '0
 * <= seq[i] < nums.len', and the binary representation of '2^seq[0] + 2^seq[1]
 * + ...' has 'k' set bits. the array product of this sequence is defined as
 * 'prod(seq) = (nums[seq[0]] * nums[seq[1]] * ...)'. return the sum of the
 * array products for all valid magical sequences. since the answer may be very
 * large, return it modulo 10^9+7. a set bit refers to a bit in the binary
 * represenation of a number that has a value of 1
 */

static int mod = 1e9 + 7;

int magicalSum(int m, int k, int *nums, int numsSize) {
  int **pow_tab = (int **)malloc(numsSize * sizeof(int *));
  for (int i = 0; i < numsSize; i++) {
    pow_tab[i] = (int *)malloc((m + 1) * sizeof(int));
    pow_tab[i][0] = 1;
    for (int j = 1; j <= m; j++) {
      long long val = (1LL * pow_tab[i][j - 1] * nums[i]) % mod;
      pow_tab[i][j] = (int)val;
    }
  }
  int **cmb = (int **)malloc((m + 1) * sizeof(int *));
  for (int i = 0; i <= m; i++) {
    cmb[i] = (int *)malloc((m + 1) * sizeof(int));
    memset(cmb[i], 0, (m + 1) * sizeof(int));
  }
  for (int i = 0; i <= m; i++) {
    cmb[i][0] = 1;
    for (int j = 1; j <= i; j++) {
      cmb[i][j] = cmb[i - 1][j - 1] + cmb[i - 1][j];
      if (cmb[i][j] >= mod)
        cmb[i][j] -= mod;
    }
  }
  int ***dp = (int ***)malloc((m + 1) * sizeof(int **));
  int ***next = (int ***)malloc((m + 1) * sizeof(int **));
  for (int i = 0; i <= m; i++) {
    dp[i] = (int **)malloc((m + 1) * sizeof(int *));
    next[i] = (int **)malloc((m + 1) * sizeof(int *));
    for (int j = 0; j <= m; j++) {
      dp[i][j] = (int *)calloc(k + 1, sizeof(int));
      next[i][j] = (int *)calloc(k + 1, sizeof(int));
    }
  }
  dp[m][0][0] = 1;
  for (int i = 0; i < numsSize; i++) {
    for (int j = 0; j <= m; j++)
      for (int c = 0; c <= m; c++)
        memset(next[j][c], 0, (k + 1) * sizeof(int));
    for (int rem = 0; rem <= m; rem++) {
      for (int carry = 0; carry <= m; carry++)
        for (int ones = 0; ones <= k; ones++) {
          int base = dp[rem][carry][ones];
          if (!base)
            continue;
          for (int c = 0; c <= rem; c++) {
            int tmp = c + carry;
            int bit = tmp & 1;
            int ones2 = ones + bit;
            if (ones2 > k)
              continue;
            int carry2 = tmp >> 1;
            int rem2 = rem - c;
            long long add = base;
            add = (add * cmb[rem][c]) % mod;
            add = (add * pow_tab[i][c]) % mod;
            int curr = next[rem2][carry2][ones2] + (int)add;
            if (curr >= mod)
              curr -= mod;
            next[rem2][carry2][ones2] = curr;
          }
        }
    }
    int ***tmp = dp;
    dp = next;
    next = tmp;
  }
  long long ans = 0;
  for (int carry = 0; carry <= m; carry++) {
    int extra = __builtin_popcount(carry);
    int need = k - extra;
    if (need >= 0 && need <= k) {
      ans += dp[0][carry][need];
      if (ans >= mod)
        ans -= mod;
    }
  }
  for (int i = 0; i < numsSize; i++)
    free(pow_tab[i]);
  free(pow_tab);
  for (int i = 0; i <= m; i++)
    free(cmb[i]);
  free(cmb);
  for (int rem = 0; rem <= m; rem++) {
    for (int carry = 0; carry <= m; carry++) {
      free(dp[rem][carry]);
      free(next[rem][carry]);
    }
    free(dp[rem]);
    free(next[rem]);
  }
  free(dp);
  free(next);
  return (int)(ans % mod);
}

int main() {
  int n1[] = {1, 10, 100, 10000, 1000000};
  int n2[] = {5, 4, 3, 2, 1};
  int n3[] = {28};
  printf("%d\n", magicalSum(5, 5, n1, ARRAY_SIZE(n1))); // expect: 991600007
  printf("%d\n", magicalSum(2, 2, n2, ARRAY_SIZE(n2))); // expect: 170
  printf("%d\n", magicalSum(1, 1, n3, ARRAY_SIZE(n3))); // expect: 28
}