Johnny and Grandmaster (Codeforces 1361B)
Tags: Greedy, sortings

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
input = sys.stdin.buffer.readline

def print(val):
    sys.stdout.write(str(val) + '\n')

def prog():
    for _ in range(int(input())):
        (n, p) = map(int, input().split())
        vals = list(map(int, input().split()))
        if p == 1:
            print(n % 2)
        else:
            vals.sort(reverse=True)
            curr_power = 0
            last = vals[0]
            too_large = False
            for a in range(n):
                k = vals[a]
                if k < last and curr_power > 0:
                    for i in range(1, last - k + 1):
                        curr_power *= p
                        if curr_power > n - a:
                            too_large = True
                            break
                    if too_large:
                        curr_power %= 1000000007
                        curr_power = curr_power * pow(p, last - i, 1000000007)
                        mod_diff = curr_power % 1000000007
                        for b in range(a, n):
                            k = vals[b]
                            mod_diff = (mod_diff - pow(p, k, 1000000007)) % 1000000007
                        print(mod_diff)
                        break
                    elif curr_power > 0:
                        curr_power -= 1
                    else:
                        curr_power += 1
                elif curr_power > 0:
                    curr_power -= 1
                else:
                    curr_power += 1
                last = k
            if not too_large:
                print(curr_power * pow(p, last, 1000000007) % 1000000007)
prog()

Solution (Python):
""""""
import os, sys, heapq as h, time
from io import BytesIO, IOBase
from types import GeneratorType
from bisect import bisect_left, bisect_right
from collections import defaultdict as dd, deque as dq, Counter as dc
import math, string
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        import os
        self.os = os
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = self.os.read(self._fd, max(self.os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = self.os.read(self._fd, max(self.os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            self.os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')

def getInt():
    return int(input())

def getStrs():
    return input().split()

def getInts():
    return list(map(int, input().split()))

def getStr():
    return input()

def listStr():
    return list(input())

def getMat(n):
    return [getInts() for _ in range(n)]

def isInt(s):
    return '0' <= s[0] <= '9'
MOD = 10 ** 9 + 7
''
T = getInt()
for _ in range(1, T + 1):
    (N, P) = getInts()
    A = getInts()
    A.sort()
    diff = 0
    prev = A[-1]
    i = N - 1
    while i >= 0:
        while not diff and i >= 1 and (A[i] == A[i - 1]):
            A.pop()
            A.pop()
            i -= 2
        if i < 0:
            break
        if diff and math.log(diff) + (prev - A[i]) * math.log(P) > math.log(i + 1):
            break
        if diff:
            diff *= pow(P, prev - A[i])
        if diff == 0:
            diff += 1
        else:
            diff -= 1
        prev = A[i]
        i -= 1
    diff *= pow(P, prev, MOD)
    while i >= 0:
        diff -= pow(P, A[i], MOD)
        i -= 1
        diff %= MOD
    print(diff)

Solution (Python):
import os
import sys
import math
from io import BytesIO, IOBase
from fractions import Fraction
import collections
from itertools import permutations
from collections import defaultdict
from collections import deque
import threading
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')

class SegmentTree:

    def __init__(self, data, default=-10 ** 6, func=lambda a, b: max(a, b)):
        """"""
        self._default = default
        self._func = func
        self._len = len(data)
        self._size = _size = 1 << (self._len - 1).bit_length()
        self.data = [default] * (2 * _size)
        self.data[_size:_size + self._len] = data
        for i in reversed(range(_size)):
            self.data[i] = func(self.data[i + i], self.data[i + i + 1])

    def __delitem__(self, idx):
        self[idx] = self._default

    def __getitem__(self, idx):
        return self.data[idx + self._size]

    def __setitem__(self, idx, value):
        idx += self._size
        self.data[idx] = value
        idx >>= 1
        while idx:
            self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1])
            idx >>= 1

    def __len__(self):
        return self._len

    def query(self, start, stop):
        if start == stop:
            return self.__getitem__(start)
        stop += 1
        start += self._size
        stop += self._size
        res = self._default
        while start < stop:
            if start & 1:
                res = self._func(res, self.data[start])
                start += 1
            if stop & 1:
                stop -= 1
                res = self._func(res, self.data[stop])
            start >>= 1
            stop >>= 1
        return res

    def __repr__(self):
        return 'SegmentTree({0})'.format(self.data)

class SegmentTree1:

    def __init__(self, data, default=10 ** 6, func=lambda a, b: min(a, b)):
        """"""
        self._default = default
        self._func = func
        self._len = len(data)
        self._size = _size = 1 << (self._len - 1).bit_length()
        self.data = [default] * (2 * _size)
        self.data[_size:_size + self._len] = data
        for i in reversed(range(_size)):
            self.data[i] = func(self.data[i + i], self.data[i + i + 1])

    def __delitem__(self, idx):
        self[idx] = self._default

    def __getitem__(self, idx):
        return self.data[idx + self._size]

    def __setitem__(self, idx, value):
        idx += self._size
        self.data[idx] = value
        idx >>= 1
        while idx:
            self.data[idx] = self._func(self.data[2 * idx], self.data[2 * idx + 1])
            idx >>= 1

    def __len__(self):
        return self._len

    def query(self, start, stop):
        if start == stop:
            return self.__getitem__(start)
        stop += 1
        start += self._size
        stop += self._size
        res = self._default
        while start < stop:
            if start & 1:
                res = self._func(res, self.data[start])
                start += 1
            if stop & 1:
                stop -= 1
                res = self._func(res, self.data[stop])
            start >>= 1
            stop >>= 1
        return res

    def __repr__(self):
        return 'SegmentTree({0})'.format(self.data)
MOD = 10 ** 9 + 7

class Factorial:

    def __init__(self, MOD):
        self.MOD = MOD
        self.factorials = [1, 1]
        self.invModulos = [0, 1]
        self.invFactorial_ = [1, 1]

    def calc(self, n):
        if n <= -1:
            print('Invalid argument to calculate n!')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        if n < len(self.factorials):
            return self.factorials[n]
        nextArr = [0] * (n + 1 - len(self.factorials))
        initialI = len(self.factorials)
        prev = self.factorials[-1]
        m = self.MOD
        for i in range(initialI, n + 1):
            prev = nextArr[i - initialI] = prev * i % m
        self.factorials += nextArr
        return self.factorials[n]

    def inv(self, n):
        if n <= -1:
            print('Invalid argument to calculate n^(-1)')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        p = self.MOD
        pi = n % p
        if pi < len(self.invModulos):
            return self.invModulos[pi]
        nextArr = [0] * (n + 1 - len(self.invModulos))
        initialI = len(self.invModulos)
        for i in range(initialI, min(p, n + 1)):
            next = -self.invModulos[p % i] * (p // i) % p
            self.invModulos.append(next)
        return self.invModulos[pi]

    def invFactorial(self, n):
        if n <= -1:
            print('Invalid argument to calculate (n^(-1))!')
            print('n must be non-negative value. But the argument was ' + str(n))
            exit()
        if n < len(self.invFactorial_):
            return self.invFactorial_[n]
        self.inv(n)
        nextArr = [0] * (n + 1 - len(self.invFactorial_))
        initialI = len(self.invFactorial_)
        prev = self.invFactorial_[-1]
        p = self.MOD
        for i in range(initialI, n + 1):
            prev = nextArr[i - initialI] = prev * self.invModulos[i % p] % p
        self.invFactorial_ += nextArr
        return self.invFactorial_[n]

class Combination:

    def __init__(self, MOD):
        self.MOD = MOD
        self.factorial = Factorial(MOD)

    def ncr(self, n, k):
        if k < 0 or n < k:
            return 0
        k = min(k, n - k)
        f = self.factorial
        return f.calc(n) * f.invFactorial(max(n - k, k)) * f.invFactorial(min(k, n - k)) % self.MOD
mod = 10 ** 9 + 7
omod = 998244353
prime = [True for i in range(50001)]
pp = []

def SieveOfEratosthenes(n=50000):
    p = 2
    while p * p <= n:
        if prime[p] == True:
            for i in range(p * p, n + 1, p):
                prime[i] = False
        p += 1
    for i in range(50001):
        if prime[i]:
            pp.append(i)
t = int(input())
for i in range(t):
    (n, p) = map(int, input().split())
    power = [int(i) for i in input().split()]
    if p == 1:
        print(n % 2)
    else:
        power.sort(reverse=True)
        ans = [0, 0]
        ok = True
        for i in range(n):
            k = power[i]
            if ans[0] == 0:
                ans = [1, k]
            elif ans[1] == k:
                ans[0] -= 1
            else:
                while ans[1] > k:
                    ans[1] -= 1
                    ans[0] *= p
                    if ans[0] >= n - i + 1:
                        ok = False
                        ind = i
                        break
                if ok == False:
                    output = ans[0] * pow(p, ans[1], mod) % mod
                    for j in range(ind, n):
                        output = (output - pow(p, power[j], mod)) % mod
                    print(output)
                    break
                else:
                    ans[0] -= 1
        if ok:
            print(ans[0] * pow(p, ans[1], mod) % mod)