3897 - Josephus Sequence

Cerința

Josephus este un matematician înrăit.

Într-o zi acesta se joacă cu primele N numere prime, când se decide să își construiască propiul său șir circular format din aceste numere. Pe prima poziție se va afla primul număr prim, adică 2, iar mai apoi se parcurge circular șirul din K în K, completându-se cu restul de numere prime, până la repartizarea tuturor.

Din nefericire lui Josephus, i-a venit somnul, așa, că vă roagă pe voi să îi construiți șirul.

Date de intrare

Fișierul de intrare josephusIN.txt conține pe prima linie numărul N și numărul K.

Date de ieșire

Fișierul de ieșire josephusOUT.txt va conține N numere naturale separate prin câte un spaţiu, reprezentând şirul lui Josephus. În cazul în care restricțiile nu sunt îndeplinite, se va afișa mesajul "Datele nu corespund restrictiilor impuse".

Restricții și precizări

• 1 ≤ N ≤ 1.000.000
• 1 ≤ K ≤ 1.000.000.000
• cel de-al 1.000.000 număr prim este 15.485.863

Exemplul 1:

josephusIN.txt

5 4

josephusOUt.txt

2 5 11 7 3

Explicație

Distribuirea se face astfel:

2 5 11 7 3 numărând din 4 în 4, începând cu primul termen, se obţine următoarea listă:

2 3 5 7 11 reprezentând primele 5 numere prime.

Exemplul 2:

josephusIN.txt

0 4

josephusOUt.txt

Datele nu corespund restrictiilor impuse

Rezolvare

import sys

def gen_primes(valMAX):
    prime = [False] * (valMAX + 1)
    primes = []
    prime[2] = True
    for i in range(3, valMAX + 1, 2):
        prime[i] = True
    
    for i in range(3, int(valMAX**0.5) + 1, 2):
        if prime[i]:
            for j in range(i*i, valMAX + 1, i + i):
                prime[j] = False
    
    primes.append(2)
    for i in range(3, valMAX + 1, 2):
        if prime[i]:
            primes.append(i)
    
    return primes

def build_aib(n):
    aib = [0] * (n + 1)
    for i in range(1, n + 1):
        aib[i] += 1
        if i + (i & -i) <= n:
            aib[i + (i & -i)] += aib[i]
    return aib

def update_add_aib(aib, n, pos, val):
    while pos <= n:
        aib[pos] += val
        pos += (pos & -pos)

def query_sum_aib(aib, pos):
    sum = 0
    while pos >= 1:
        sum += aib[pos]
        pos -= (pos & -pos)
    return sum

def query_range_sum_aib(aib, le, ri):
    return query_sum_aib(aib, ri) - query_sum_aib(aib, le - 1)

def bs_aib(aib, n, val, logn):
    sum = 0
    pos = 0
    for i in range(logn, -1, -1):
        if pos + (1 << i) <= n and sum + aib[pos + (1 << i)] < val:
            sum += aib[pos + (1 << i)]
            pos += (1 << i)
    return pos + 1

def put_int(x):
    sys.stdout.write(str(x) + ' ')

def check_restrictions(n, k):
    return 1 <= n <= 1000000 and 1 <= k <= 1000000000

def main():
    with open("josephusIN.txt", "r") as fin:
        n, k = map(int, fin.read().split())

    with open("josephusOUT.txt", "w") as fout:
        if not check_restrictions(n, k):
            fout.write("Datele nu corespund restrictiilor impuse\n")
            return
        
        valMAX = 15485863
        primes = gen_primes(valMAX)
        logn = 31 - (n).bit_length() + 1
        aib = build_aib(n)
        
        v = [0] * (n + 1)
        v[1] = 2
        update_add_aib(aib, n, 1, -1)
        
        st = 1
        for i in range(2, n + 1):
            ezk = (k - 1) % query_sum_aib(aib, n) + 1
            
            if query_range_sum_aib(aib, st, n) >= ezk:
                st = bs_aib(aib, n, ezk + query_sum_aib(aib, st - 1), logn)
            else:
                st = bs_aib(aib, n, ezk - query_range_sum_aib(aib, st, n), logn)
            
            v[st] = primes[i - 1]
            update_add_aib(aib, n, st, -1)
        
        for i in range(1, n + 1):
            fout.write(f"{v[i]} ")
        fout.write("\n")

if __name__ == "__main__":
    main()