HackerRank Combinatorics – Count Palindromes

Challenge: Count Palindromes

Subdomeniu: Combinatorics (combinatorics)

Scor cont: 140.0 / 140

Submission status: Accepted

Submission score: 1.0

Submission ID: 464834583

Limbaj: cpp14

Link challenge: https://www.hackerrank.com/challenges/count-palindromes/problem

Cerinta

A string is made of only lowercase latin letters (a,b,c,d,.....,z). Can you find the length of the lexicographically smallest string such that it has exactly $K$ sub-strings, each of which are palindromes?

Input Format

The first line of input contains single integer $T$ - the number of testcases.  
T lines follow, each containing the integer $K$.

Output Format

Output exactly $T$ lines. Each line should contain single integer - the length of the lexicographically smallest string.

Constraints

* $1 \leq T \leq 100$  
* $1 \leq K \leq 10^{12}$

Cod sursa

#include <bits/stdc++.h>
using namespace std;

// Count Palindromes - HackerRank
// Given K, output the length of the lexicographically smallest string with exactly K palindromic substrings.

static inline __int128 sum_arith(long long first, long long len){
    // first + (first+1) + ... + (first+len-1)
    __int128 L = len;
    __int128 a = first;
    return L * (2*a + (L-1)) / 2;
}

static long long max_len_run(long long k, long long firstCost){
    // maximum L >= 1 such that sum_arith(firstCost, L) <= k
    // Solve: L^2 + (2*firstCost-1)L - 2k <= 0
    long double b = (long double)(2*firstCost - 1);
    long double disc = b*b + (long double)8.0L * (long double)k;
    long double root = sqrt(disc);

    long long L = (long long)floor((-b + root) / 2.0L);
    if(L < 1) L = 1;

    while(sum_arith(firstCost, L) > (__int128)k) L--;
    while(sum_arith(firstCost, L+1) <= (__int128)k) L++;
    return L;
}

static long long solve_one(long long k){
    long long ans = 0;
    long long costA = 1; // cost to append 'a'
    long long costB = 1; // cost to append 'b' (1 only before first 'b', then becomes 2)

    while(k > 0){
        if(costA <= k){
            long long L = max_len_run(k, costA);
            __int128 s = sum_arith(costA, L);
            k -= (long long)s;
            ans += L;
            costA += L;
        } else if(costB <= k){
            k -= costB;
            ans += 1;
            if(costB == 1){
                costA = 2;
                costB = 2;
            } else { // costB == 2
                costA = 3;
            }
        } else {
            // only possible when k == 1: append 'c'
            ans += 1;
            k -= 1;
        }
    }
    return ans;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int T;
    cin >> T;
    while(T--){
        long long K;
        cin >> K;
        cout << solve_one(K) << "\n";
    }
    return 0;
}