Problem 822: Square the Smallest

Problem 822 source code

C++

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <queue>
#include <vector>

using u64 = std::uint64_t;
using u128 = unsigned __int128;

static constexpr u64 kMod = 1'234'567'891ULL;

struct Node {
    long double logv;
    u64 modv;
    int id;
};

struct MinCmp {
    bool operator()(const Node& a, const Node& b) const {
        if (a.logv != b.logv) return a.logv > b.logv;
        return a.id > b.id;
    }
};

static u64 mod_pow(u64 base, u64 exp, u64 mod) {
    u64 res = 1 % mod;
    base %= mod;
    while (exp > 0) {
        if (exp & 1ULL) {
            res = static_cast<u64>((static_cast<u128>(res) * base) % mod);
        }
        base = static_cast<u64>((static_cast<u128>(base) * base) % mod);
        exp >>= 1ULL;
    }
    return res;
}

static u64 S(u64 n, u64 m) {
    const int len = static_cast<int>(n - 1);

    std::priority_queue<Node, std::vector<Node>, MinCmp> pq;
    std::vector<Node> cur(len);
    long double max_log = 0.0L;

    for (int i = 0; i < len; ++i) {
        const u64 v = static_cast<u64>(i + 2);
        cur[i] = {std::log(static_cast<long double>(v)), v % kMod, i};
        pq.push(cur[i]);
        if (cur[i].logv > max_log) {
            max_log = cur[i].logv;
        }
    }

    u64 steps = 0;
    while (steps < m) {
        Node x = pq.top();
        if (2.0L * x.logv > max_log) {
            break;
        }
        pq.pop();
        x.logv *= 2.0L;
        x.modv = static_cast<u64>((static_cast<u128>(x.modv) * x.modv) % kMod);
        pq.push(x);
        if (x.logv > max_log) {
            max_log = x.logv;
        }
        ++steps;
    }

    std::vector<Node> arr;
    arr.reserve(len);
    while (!pq.empty()) {
        arr.push_back(pq.top());
        pq.pop();
    }

    if (steps == m) {
        u64 ans = 0;
        for (const auto& x : arr) {
            ans += x.modv;
            if (ans >= kMod) {
                ans -= kMod;
            }
        }
        return ans;
    }

    std::sort(arr.begin(), arr.end(), [](const Node& a, const Node& b) {
        if (a.logv != b.logv) return a.logv < b.logv;
        return a.id < b.id;
    });

    const u64 rem = m - steps;
    const u64 q = rem / static_cast<u64>(len);
    const u64 r = rem % static_cast<u64>(len);

    const u64 exp_q = mod_pow(2, q, kMod - 1);
    const u64 exp_q1 = static_cast<u64>((static_cast<u128>(exp_q) * 2ULL) % (kMod - 1));

    u64 ans = 0;
    for (int i = 0; i < len; ++i) {
        const u64 e = (static_cast<u64>(i) < r) ? exp_q1 : exp_q;
        const u64 add = mod_pow(arr[i].modv, e, kMod);
        ans += add;
        if (ans >= kMod) {
            ans -= kMod;
        }
    }

    return ans;
}

int main() {
    assert(S(5, 3) == 34);
    assert(S(10, 100) == 845339386ULL);
    std::cout << S(10'000, 10'000'000'000'000'000ULL) << '\n';
    return 0;
}

Python

import math
import heapq

kMod = 1234567891

class Node:
    def __init__(self, logv, modv, id):
        self.logv = logv
        self.modv = modv
        self.id = id
        
    def __lt__(self, other):
        if self.logv != other.logv:
            return self.logv < other.logv
        return self.id < other.id

def S(n, m):
    len_n = n - 1
    pq = []
    max_log = 0.0
    
    for i in range(len_n):
        v = i + 2
        logv = math.log(v)
        node = Node(logv, v % kMod, i)
        heapq.heappush(pq, node)
        if logv > max_log:
            max_log = logv
            
    steps = 0
    while steps < m:
        x = pq[0]
        if 2.0 * x.logv > max_log:
            break
        heapq.heappop(pq)
        x.logv *= 2.0
        x.modv = (x.modv * x.modv) % kMod
        heapq.heappush(pq, x)
        if x.logv > max_log:
            max_log = x.logv
        steps += 1
        
    if steps == m:
        ans = 0
        for x in pq:
            ans = (ans + x.modv) % kMod
        return ans
        
    arr = list(pq)
    arr.sort(key=lambda x: (x.logv, x.id))
    
    rem = m - steps
    q = rem // len_n
    r = rem % len_n
    
    exp_q = pow(2, q, kMod - 1)
    exp_q1 = (exp_q * 2) % (kMod - 1)
    
    ans = 0
    for i in range(len_n):
        e = exp_q1 if i < r else exp_q
        add = pow(arr[i].modv, e, kMod)
        ans = (ans + add) % kMod
        
    return ans

def solve():
    ans = S(10000, 10000000000000000)
    return str(ans)

if __name__ == "__main__":
    print(solve())

Java

import java.util.PriorityQueue;
import java.util.Comparator;
import java.util.ArrayList;
import java.util.Collections;

public class Euler822 {

    static final long kMod = 1234567891L;

    static class Node {
        double logv;
        long modv;
        int id;

        Node(double logv, long modv, int id) {
            this.logv = logv;
            this.modv = modv;
            this.id = id;
        }
    }

    static long modPow(long base, long exp, long mod) {
        long res = 1 % mod;
        base %= mod;
        while (exp > 0) {
            if ((exp & 1L) == 1L) {
                res = (res * base) % mod;
            }
            base = (base * base) % mod;
            exp >>= 1L;
        }
        return res;
    }

    static long S(long n, long m) {
        int len = (int) (n - 1);

        PriorityQueue<Node> pq = new PriorityQueue<>(new Comparator<Node>() {
            @Override
            public int compare(Node a, Node b) {
                if (a.logv != b.logv) {
                    return Double.compare(a.logv, b.logv);
                }
                return Integer.compare(a.id, b.id);
            }
        });

        double maxLog = 0.0;

        for (int i = 0; i < len; ++i) {
            long v = i + 2;
            Node node = new Node(Math.log(v), v % kMod, i);
            pq.add(node);
            if (node.logv > maxLog) {
                maxLog = node.logv;
            }
        }

        long steps = 0;
        while (steps < m) {
            Node x = pq.peek();
            if (2.0 * x.logv > maxLog) {
                break;
            }
            pq.poll();
            x.logv *= 2.0;
            x.modv = (x.modv * x.modv) % kMod;
            pq.add(x);
            if (x.logv > maxLog) {
                maxLog = x.logv;
            }
            ++steps;
        }

        ArrayList<Node> arr = new ArrayList<>(pq);

        if (steps == m) {
            long ans = 0;
            for (Node x : arr) {
                ans += x.modv;
                if (ans >= kMod) {
                    ans -= kMod;
                }
            }
            return ans;
        }

        Collections.sort(arr, new Comparator<Node>() {
            @Override
            public int compare(Node a, Node b) {
                if (a.logv != b.logv) {
                    return Double.compare(a.logv, b.logv);
                }
                return Integer.compare(a.id, b.id);
            }
        });

        long rem = m - steps;
        long q = rem / len;
        long r = rem % len;

        long expQ = modPow(2, q, kMod - 1);
        long expQ1 = (expQ * 2L) % (kMod - 1);

        long ans = 0;
        for (int i = 0; i < len; ++i) {
            long e = (i < r) ? expQ1 : expQ;
            long add = modPow(arr.get(i).modv, e, kMod);
            ans += add;
            if (ans >= kMod) {
                ans -= kMod;
            }
        }

        return ans;
    }

    public static String solve() {
        return Long.toString(S(10000, 10000000000000000L));
    }

    public static void main(String[] args) {
        System.out.println(solve());
    }
}