Johnny and Grandmaster (Codeforces 1362E)
Tags: Greedy, number theory

Johnny has just found the new, great tutorial: "How to become a grandmaster?". The tutorial tells many strange and unexpected for Johnny things, such as you have to be patient or that very important is solving many harder and harder problems. The boy has found an online judge with tasks divided by topics they cover. He has picked $$$p^{k_i}$$$ problems from $$$i$$$-th category ($$$p$$$ is his favorite number). He wants to solve them in two weeks (the patience condition is too hard for Johnny, so for simplicity, he looks only at easy tasks, which can be solved in such a period). Now our future grandmaster has to decide which topics to cover first and which the second week. Help him assign topics in such a way, that workload is balanced.Formally, given $$$n$$$ numbers $$$p^{k_i}$$$, the boy wants to divide them into two disjoint sets, minimizing the absolute difference between sums of numbers in each set. Find the minimal absolute difference. Output the result modulo $$$10^{9}+7$$$.

Input: Input consists of multiple test cases. The first line contains one integer $$$t$$$ $$$(1 \leq t \leq 10^5)$$$ — the number of test cases. Each test case is described as follows:The first line contains two integers $$$n$$$ and $$$p$$$ $$$(1 \leq n, p \leq 10^6)$$$. The second line contains $$$n$$$ integers $$$k_i$$$ $$$(0 \leq k_i \leq 10^6)$$$.The sum of $$$n$$$ over all test cases doesn't exceed $$$10^6$$$.

Output: Output one integer — the reminder of division the answer by $$$1\,000\,000\,007$$$.

Note: You have to minimize the difference, not it's remainder. For example, if the minimum difference is equal to $$$2$$$, but there is also a distribution where the difference is $$$10^9 + 8$$$, then the answer is $$$2$$$, not $$$1$$$.In the first test case of the example, there're the following numbers: $$$4$$$, $$$8$$$, $$$16$$$, $$$16$$$, and $$$8$$$. We can divide them into such two sets: $$${4, 8, 16}$$$ and $$${8, 16}$$$. Then the difference between the sums of numbers in sets would be $$$4$$$.

Sample Inputs:
4
5 2
2 3 4 4 3
3 1
2 10 1000
4 5
0 1 1 100
1 8
89

Sample Outputs:
4
1
146981438
747093407

Solution (Python):
import sys
from collections import deque, Counter
input = sys.stdin.buffer.readline
MOD = 1000000007
T = int(input())
for _ in range(T):
    (n, p) = map(int, input().split())
    ls = list(map(int, input().split()))
    ls.sort(reverse=True)
    if p == 1:
        print(1 if n % 2 else 0)
        continue
    (cp, cc) = (0, 0)
    (ok, pos) = (1, 0)
    for (i, u) in enumerate(ls):
        if cc == 0:
            (cp, cc) = (u, 1)
        else:
            (tp, tc) = (cp, cc)
            while tp > u and tc <= n:
                tp -= 1
                tc *= p
            if tp == u:
                (cp, cc) = (tp, tc - 1)
            else:
                (ok, pos) = (0, i)
                break
    if ok:
        print(pow(p, cp, MOD) * cc % MOD)
    else:
        a = pow(p, cp, MOD) * cc % MOD
        b = 0
        for u in ls[pos:]:
            b += pow(p, u, MOD)
            b %= MOD
        print((a - b) % MOD)

Solution (Python):
import sys
from collections import deque, Counter
input = sys.stdin.buffer.readline
MOD = 1000000007
T = int(input())
for _ in range(T):
    (n, p) = map(int, input().split())
    ls = list(map(int, input().split()))
    ls.sort(reverse=True)
    (cp, cc) = (0, 0)
    (ok, pos) = (1, 0)
    for (i, u) in enumerate(ls):
        if cc == 0:
            (cp, cc) = (u, 1)
        else:
            (tp, tc) = (cp, cc)
            if p == 1:
                tp = u
            else:
                while tp > u and tc <= n:
                    tp -= 1
                    tc *= p
            if tp == u:
                (cp, cc) = (tp, tc - 1)
            else:
                (ok, pos) = (0, i)
                break
    if ok:
        print(pow(p, cp, MOD) * cc % MOD)
    else:
        a = pow(p, cp, MOD) * cc % MOD
        b = 0
        for u in ls[pos:]:
            b += pow(p, u, MOD)
            b %= MOD
        print((a - b) % MOD)

Solution (Python):
import sys, math
import io, os

def data():
    return sys.stdin.buffer.readline().strip()

def mdata():
    return list(map(int, data().split()))

def outl(var):
    sys.stdout.write(' '.join(map(str, var)) + '\n')

def out(var):
    sys.stdout.write(str(var) + '\n')
mod = int(1000000000.0) + 7

def main():
    for t in range(int(data())):
        (n, p) = mdata()
        K = mdata()
        (a, b) = (0, 0)
        if p == 1:
            out(n % 2)
            continue
        flag = 0
        K.sort(reverse=True)
        for i in range(n):
            k = K[i]
            if a == 0:
                b = k
                a = 1
            elif b == k:
                a -= 1
            else:
                while b > k:
                    a *= p
                    b -= 1
                    if a > n - i - 1:
                        flag = 1
                        i -= 1
                        break
                if flag == 1:
                    break
                a -= 1
        ans = (a * pow(p, b, mod) - sum((pow(p, K[j], mod) for j in range(i + 1, n))) + mod) % mod
        out(ans)
if __name__ == '__main__':
    main()