Sum of the sum of all possible subsets raised to power k

Question

You are given an array A of non-negative integers. You can pick any non-empty subset, S from the array A. The score of a subset S is the sum of the elements in S raised to the power of K, i.e. for a subset S={s₁,s₂,…,s_m}, the score of S is (s₁+s₂+…,s_m)^K. Output the sum of scores over all possible non-empty subsets of A modulo 998244353.

Input

The first line consists of two integers N (1 ≤ N ≤ 1000) and K (1 ≤ K ≤ 300), Then N integers follow: a₁,a₂…,a_n (1 ≤ a_i ≤ 10⁹)

Examples

Input:

3 2
1 2 4

Output:

140

Note

There are 7 possible non empty subsets:

{1}, 1²=1

{2}, 2²=4

{4}, 4²=16

{1,2}, 3²=9

{1,4}, 5²=25

{2,4}, 6²=36

{1,2,4}, 7²=49

The total of all of them is 140.

Test Case 1

Input:

200 33
586675996 834662638 582661564 801031252 368588510 481535057 299143636 154810868 748560222 444723881 594265644 271475209 483236369 825333010 838212251 941340273 181680951 475018532 725610690 980879808 290253341 255300094 222425864 305836905 309362719 673795903 526752784 272885696 498600892 541576513 293569958 586502788 546400101 932614930 922795145 738713347 627678736 854811027 599387102 172776156 368969197 453280296 317871960 878976720 860453033 510602497 926525753 419146537 830612914 980672268 468005306 166335891 911567226 922140110 434524807 424093053 590300567 371183659 721172742 600311592 210722614 317227418 793158971 401617057 449556863 501677314 892013630 839176022 524835881 940550248 839561467 764011513 989141930 192333644 484181005 785400847 674841029 209630134 964257438 941649848 836682055 640053044 493288668 662886952 142491554 659229917 884165257 346190427 135295371 636702497 710634145 170177475 893751212 553144583 498928648 173846992 274295496 676101828 165556726 499077801 920844845 446218268 371964528 518465222 701401156 687218925 309703289 579304688 823881043 802733954 252512484 579282959 844367812 695500833 748122407 825455122 386300070 502934300 179547526 934792800 783118124 441360402 356559179 327130654 998543269 475814368 713846691 788491645 617178564 384577123 166216557 684039923 651593797 821711021 704057885 243187642 943485338 885996539 369725079 504573341 230128629 304769483 940749404 725863600 460266942 378887332 700635348 115095533 583723838 352744521 560578526 984014856 678045467 988770077 470484548 380200831 349213692 849092703 695155761 125390576 970541165 842717904 202532794 848261409 892205344 324076314 401291603 958139378 960397596 669808951 333662236 895479878 866907097 938228867 851281251 333564143 552098658 371049021 356458887 441896629 907071903 336678292 857795316 302950189 898249304 663033041 993067004 494024812 140930107 733086617 107029928 909953444 616400129 799844902 322302802 644050662 253833551 636316553 633766994 300303396

Scoring

The goal for this challenge is the produce the fastest algorithm (i.e, the algorithm with the smallest asymptotic complexity), and as such you should include an short analysis of your algorithm alongside your code.

Answer 1

Python, \$\mathcal O(NK^2)\$

from math import comb
from functools import lru_cache

M = 998244353

rec=lru_cache(None)(
lambda n,a,*A:(pow(a,n,M)+(len(A)and rec(n,*A)+sum(rec(n-j,*A)*pow(a,j,M)*comb(n,j)for j in range(n+1))))%M)

Attempt This Online!

\$\mathcal O(NK^2)\$ (length of list times exponent squared) assuming that binomial coefficients and powers are O(1) (which is not true for general random access but can be achieved when the values to compute are nearby because then one can be computed from a neighbouring cheaply, cf. comment by @Bubbler. To be clear, I did not implement this optimisation because I didn't find it worth the clutter.) and the overhead of memoisation is negligible.

Does the large testcase in a second or so.

Additional testcases taken from @bsoelch.

How?

Let \$f(K,S)\$ denote the desired quantity for exponent K and set (or multiset) S. I write \$2^S\$ for the power set. The recursion formula is obtained by choosing an element \$a\in S\$ and separating the subsets into those that contain \$a\$ and those that do not.

$$f(K,S) = \sum_{C \in 2^S}\left(\sum_{x \in C}x\right)^K = \sum_{C \in 2^{S-\{a\}}}\left(\sum_{x \in C}x\right)^K + \sum_{C \in 2^{S-\{a\}}}\left(a+\sum_{x \in C} x\right)^K \\ = f(K,S-\{a\}) + \sum_{C \in 2^{S-\{a\}}}\left(a+\sum_{x \in C} x\right)^K \\ = f(K,S-\{a\}) + \sum_{C \in 2^{S-\{a\}}}\sum_{j=0}^K \begin{pmatrix} K\\j \end{pmatrix}a^j \left(\sum_{x \in C} x\right)^{K-j} \\ = f(K,S-\{a\}) + \sum_{j=0}^K \begin{pmatrix} K\\j \end{pmatrix}a^j \sum_{C \in 2^{S-\{a\}}} \left(\sum_{x \in C} x\right)^{K-j} \\ =f (K,S-\{a\}) + \sum_{j=0}^K \begin{pmatrix} K\\j \end{pmatrix}a^j f(K-j,S-\{a\}) $$

Answer 2

Python, \$O(N K \log K)\$

from functools import lru_cache
from sympy.discrete.convolutions import convolution_ntt

MOD = 998244353


@lru_cache(maxsize=None)
def factorial(n):
    return n * factorial(n - 1) % MOD if n else 1


@lru_cache(maxsize=None)
def inverse_factorial(n):
    return pow(factorial(n), -1, MOD)


def r(A, K):
    if len(A) == 0:
        return [1] + [0] * K
    p1 = r(A[1:], K)
    p2 = [pow(A[0], i, MOD) * inverse_factorial(i) % MOD for i in range(K + 1)]
    p = convolution_ntt(p1, p2, MOD)[:K + 1]
    rv = [(a + b) % MOD for a, b in zip(p1, p)]
    return rv


def f(A, K):
    return r(A, K)[-1] * factorial(K) % MOD

Attempt This Online!

Continuing from @loopy walt's recursion, we have $$f(K, S) = f(K, S - \{a\}) + \sum_{j=0}^K {k! \frac{a^j}{j!} \frac{f(K-j, S-\{a\})}{(K-j)!}} $$ so we look at the EGF \$f(S) = \sum_{k=0}^{\infty} {f(k, S) \frac{x^k}{k!}}\$ and we have \$f(S) = (1 + e^{a x}) f(S - \{a\})\$, so we evaluate this using NTT, cutting off the terms greater than \$x^K\$. We do \$N\$ convolutions on arrays of length \$K\$, so the time complexity is \$O(N K \log K)\$.

Answer 3

Python, 428 bytes, \$O(m^{k+1})\$

Ignores the input format requirements, f takes input (l,k) with l being the array and k being the power.
If it is really necessary I can add a parsing function (see ¹), to the code, but this will not impact the total complexity and make it harder to check multiple cases at once

from itertools import product
from math import comb 

M=998244353

def prod(l): ## product over the list l
  P=1
  for x in l:
   P*=x
   if M:
     P%=M  ## only need result modulo M, so take mod M to make numbers smaller
  return P

def f(l,k):
  n=len(l)
  S=0
  for P in product(range(n),repeat=k):
   S+=prod(l[i] for i in P)*pow(2,(n-len(set(P))),M)
   if M:
     S%=M  ## only need result modulo M, so take mod M to make numbers smaller
  return S

Attempt This Online! (runs for ~30 seconds)

Explanation

Idea: using the generalized binomial formula split the term for each subset into summands and count how often each combination of elements appears in the final sum.

Example:

{x,y,z}, 2 ->
0       
x       -> x²
  y     ->    y²
x,y     -> x² y²    2xy
    z   ->       z²
x,  z   -> x²    z²     2xz
  y,z   ->    y² z²         2yz
x,y,z   -> x² y² z² 2xy 2xz 2yz

Each term appears (counting different orderings multiple times) \$2^{len(l)-|\{\text{different elements in term}\}|}\$ times (the number of subsets a term appears in halves with each unique letter used in the term). So we can get the total sum by summing up the term multiplied by their frequencies.

Running time

Let \$m\$ be the length of the list and \$k\$ be the exponent. The program checks \$O(m^{k})\$ different terms, for each of terms the computation can be done in \$O(m)\$ elementary calculations. As each elementary computation modulo 998244353 can be done in \$O(1)\$ this gives a total running time of \$O(m^{k+1})\$

Combining all terms using the same elements and multiplying by the number of permutations would not reduce the running time as there are still \${m \choose k} \in O(m^k)\$ different terms

---

¹parsing function:

def parse():
  _,k=[*map(int,input().split())]
  l=[*map(int,input().split())]
  return l,k

Attempt This Online!

Answer 4

Rust, \$\mathcal{O}(NK^2)\$

A port of @loopy walt's Python answer in Rust.

---

src/main.rs

use cached::proc_macro::cached;
use num::integer::binomial;
// use num::pow::pow;

const M: i64 = 998244353;

fn pow_mod(base: i64, exp: usize, m: i64) -> i64 {
    let mut result = 1;
    let mut base = base % m;
    let mut exp = exp;

    while exp > 0 {
        if exp % 2 == 1 {
            result = (result * base) % m;
        }
        exp >>= 1;
        base = (base * base) % m;
    }
    (result+m)%m
}

#[cached]
fn rec(n: u32, a: Vec<i64>) -> i64 {
    if a.is_empty() {
        return 0;
    }
    let mut result: i64 = pow_mod(a[0], n as usize, M);
    if a.len() > 1 {
        let mut rest = a.clone();
        rest.remove(0);
        result = ( result+ rec(n, rest))%M;
        for j in 0..=n {
            let mut rest_j = a.clone();
            rest_j.remove(0);
            result = (result + (rec(n - j, rest_j) * pow_mod(a[0] as i64, j as usize, M)%M * binomial(n, j) as i64 %M))%M;
        }
    }
    (result + M)% M
}

fn main() {
    println!("{}", rec(2, vec![1, 2, 4]));
    println!("{}", rec(2, vec![1, 2, 3]));
    println!("{}", rec(3, vec![1, 2, 3, 4]));
    println!("{}", rec(9, vec![3, 1, 4, 1, 5]));
    println!("{}", rec(33, vec![586675996,834662638,582661564,801031252,368588510,481535057,299143636,154810868,748560222,444723881,594265644,271475209,483236369,825333010,838212251,941340273,181680951,475018532,725610690,980879808,290253341,255300094,222425864,305836905,309362719,673795903,526752784,272885696,498600892,541576513,293569958,586502788,546400101,932614930,922795145,738713347,627678736,854811027,599387102,172776156,368969197,453280296,317871960,878976720,860453033,510602497,926525753,419146537,830612914,980672268,468005306,166335891,911567226,922140110,434524807,424093053,590300567,371183659,721172742,600311592,210722614,317227418,793158971,401617057,449556863,501677314,892013630,839176022,524835881,940550248,839561467,764011513,989141930,192333644,484181005,785400847,674841029,209630134,964257438,941649848,836682055,640053044,493288668,662886952,142491554,659229917,884165257,346190427,135295371,636702497,710634145,170177475,893751212,553144583,498928648,173846992,274295496,676101828,165556726,499077801,920844845,446218268,371964528,518465222,701401156,687218925,309703289,579304688,823881043,802733954,252512484,579282959,844367812,695500833,748122407,825455122,386300070,502934300,179547526,934792800,783118124,441360402,356559179,327130654,998543269,475814368,713846691,788491645,617178564,384577123,166216557,684039923,651593797,821711021,704057885,243187642,943485338,885996539,369725079,504573341,230128629,304769483,940749404,725863600,460266942,378887332,700635348,115095533,583723838,352744521,560578526,984014856,678045467,988770077,470484548,380200831,349213692,849092703,695155761,125390576,970541165,842717904,202532794,848261409,892205344,324076314,401291603,958139378,960397596,669808951,333662236,895479878,866907097,938228867,851281251,333564143,552098658,371049021,356458887,441896629,907071903,336678292,857795316,302950189,898249304,663033041,993067004,494024812,140930107,733086617,107029928,909953444,616400129,799844902,322302802,644050662,253833551,636316553,633766994,300303396]));
}

Cargo.toml

[package]
name = "rust_hello"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
cached = "0.26.2"
num = "0.4.0"

[profile.release]
lto = "y"
panic = "abort"
rand = "0.8.0"

Build and run

$ cargo build --release
$ cargo run --release
140
100
3800
554473082
856382558