CODEWITHSUNDEEP
algorithm
dorado
dorado

Reputation: 1525

Count number of subsequences with given k modulo sum

Given an array a of n integers, count how many subsequences (non-consecutive as well) have sum % k = 0:

1 <= k < 100
1 <= n <= 10^6
1 <= a[i] <= 1000

An O(n^2) solution is easily possible, however a faster way O(n log n) or O(n) is needed.

Upvotes: 7

Views: 2345

Answers (3)

גלעד ברקן
גלעד ברקן

Reputation: 23955

Traverse a and count a[i] mod k; there ought to be k such counts.

Recurse and memoize over the distinct partitions of k, 2*k, 3*k...etc. with parts less than or equal to k, adding the products of the appropriate counts.

For example, if k were 10, some of the partitions would be 1+2+7 and 1+2+3+4; but while memoizing, we would only need to calculate once how many pairs mod k in the array produce (1 + 2).

For example, k = 5, a = {1,4,2,3,5,6}:

counts of a[i] mod k: {1,2,1,1,1}

products of distinct partitions of k:
  5   => 1
  4,1 => 2
  3,2 => 1

products of distinct partitions of 2 * k with parts <= k:
  5,4,1   => 2
  5,3,2   => 1
  4,1,3,2 => 2

products of distinct partitions of 3 * k with parts <= k:
  5,4,1,3,2 => 2

answer = 11

  {1,4} {4,6} {2,3} {5}
  {1,4,2,3} {1,4,5} {4,6,2,3} {4,6,5} {2,3,5}
  {1,4,2,3,5} {4,6,2,3,5}

Upvotes: 0

David Eisenstat
David Eisenstat

Reputation: 65506

There's an O(n + k^2 lg n)-time algorithm. Compute a histogram c(0), c(1), ..., c(k-1) of the input array mod k (i.e., there are c(r) elements that are r mod k). Then compute

  k-1
product (1 + x^r)^c(r) mod (1 - x^k)
  r=0

as follows, where the constant term of the reduced polynomial is the answer.

Rather than evaluate each factor with a fast exponentiation method and then multiply, we turn things inside out. If all c(r) are zero, then the answer is 1. Otherwise, recursively evaluate

      k-1
P = product (1 + x^r)^(floor(c(r)/2)) mod (1 - x^k).
      r=0

and then compute

      k-1
Q = product (1 + x^r)^(c(r) - 2 floor(c(r)/2)) mod (1 - x^k),
      r=0

in time O(k^2) for the latter computation by exploiting the sparsity of the factors. The result is P^2 Q mod (1 - x^k), computed in time O(k^2) via naive convolution.

Upvotes: 0

IVlad
IVlad

Reputation: 43517

This is the subset sum problem.

A simple solution is this:

s = 0
dp[x] = how many subsequences we can build with sum x 
dp[0] = 1, 0 elsewhere
for i = 1 to n:
    s += a[i]
    for j = s down to a[i]:
        dp[j] = dp[j] + dp[j - a[i]]

Then you can simply return the sum of all dp[x] such that x % k == 0. This has a high complexity though: about O(n*S), where S is the sum of all of your elements. The dp array must also have size S, which you probably can't even afford to declare for your constraints.

A better solution is to not iterate over sums larger than or equal to k in the first place. To do this, we will use 2 dp arrays:

dp1, dp2 = arrays of size k
dp1[0] = dp2[0] = 1, 0 elsewhere
for i = 1 to n:
    mod_elem = a[i] % k
    for j = 0 to k - 1:
        dp2[j] = dp2[j] + dp1[(j - mod_elem + k) % k]

    copy dp2 into dp1

return dp1[0]

Whose complexity is O(n*k), and is optimal for this problem.

Upvotes: 3

Related Questions