Test Case Pass: Failed on Test Case 9 ❌

blog_gemini_2.5_coding_cf_zebra_like_numbers_question_response.cpp

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>

using namespace std;

// Define long long type for large numbers
typedef long long ll;

// Use a map for memoization. The key is a pair {N, k}, representing the state (upper bound N, target zebra value k).
// The value is the computed count for this state.
map<pair<ll, ll>, ll> memo;

// Store precomputed zebra numbers in a vector.
vector<ll> Z;

// Function to precompute zebra numbers.
// Zebra numbers are of the form 1, 101, 10101, ... in binary.
// In decimal: 1, 5, 21, 85, ...
// The formula is Z_m = (4^k - 1) / 3 where the binary representation has length m=2k-1.
// We precompute them up to a value slightly larger than 10^18.
void precompute_zebra() {
    ll current = 1;
    // Always add 1, the first zebra number.
    Z.push_back(current); 
    ll power_of_4 = 1;
    
    // Loop to generate subsequent zebra numbers.
    while (true) { 
        // Check for potential overflow before multiplying power_of_4 by 4.
        // __LONG_LONG_MAX__ is the maximum value for long long.
        if (power_of_4 > __LONG_LONG_MAX__ / 4) break;
        power_of_4 *= 4;

        // Check for potential overflow before adding power_of_4 to current.
        if (current > __LONG_LONG_MAX__ - power_of_4) break;
        current += power_of_4;
        
        // Define a practical upper limit slightly above 10^18 to avoid generating excessively large numbers.
        // 10^18 + 10^17 is a heuristic margin.
        // Using 1e18 notation requires casting or suffix LL. Use explicit values.
        ll limit = 1000000000000000000LL + 100000000000000000LL; 
        if (current > limit) break; 
        
        // Add the newly computed zebra number to the vector.
        Z.push_back(current);
    }
}

// Recursive function to count numbers x in the interval [1, N] such that their zebra value is exactly k.
ll count(ll N, ll k) {
    // Base cases:
    // If k is non-positive, no positive integer has such zebra value.
    if (k <= 0) return 0;
    // If N is non-positive, the interval [1, N] is empty.
    if (N <= 0) return 0; 
    
    // Optimization: Based on empirical observation and problem structure, it's likely that
    // the zebra value p(x) for x <= 10^18 is relatively small. A heuristic upper bound
    // (e.g., 70) can significantly prune the search space for very large k. If k exceeds this bound,
    // we assume the count is 0. This bound might need adjustment if contradicted by test cases.
    if (k > 70) return 0; 

    // Check if the result for state (N, k) is already computed and stored in the memoization table.
    pair<ll, ll> state = {N, k};
    if (memo.count(state)) return memo[state];
    
    // Find the largest zebra number Z_m <= N.
    // `upper_bound` returns an iterator to the first element in Z strictly greater than N.
    auto it = upper_bound(Z.begin(), Z.end(), N);
    
    // If N is smaller than the smallest zebra number (1), then no number in [1, N] can be formed.
    // The interval [1, N] would be empty or contain numbers less than 1.
    if (it == Z.begin()) { 
         return memo[state] = 0; 
    }
    // Decrement the iterator to point to the largest element Z_m <= N.
    it--; 
    ll Z_m = *it; // Store the value of Z_m.
    
    ll res; // Variable to store the result for state (N, k).
    
    // Apply the recursive formula based on k.
    // This formula relies on the property p(x) = p(x - Z_m) + 1, where Z_m is the largest zebra number <= x.
    // Count(N, k) = Count(Z_m-1, k) + # {x | Z_m <= x <= N and p(x) = k}
    // The second term is equivalent to # {y | 0 <= y <= N-Z_m and p(y) = k-1}
    
    if (k == 1) {
        // Case k=1: We are counting zebra numbers <= N.
        // count(Z_m - 1, 1) counts zebra numbers strictly less than Z_m.
        // We add 1 to include Z_m itself, since Z_m <= N.
        // The count of numbers y with p(y) = 0 is 1 (only y=0).
        res = count(Z_m - 1, 1) + 1; 
    } else {
        // Case k > 1: Apply the main recurrence relation.
        // Count numbers in [1, Z_m-1] with value k PLUS count numbers in [Z_m, N] with value k.
        // The count for [Z_m, N] translates to count numbers y in [1, N - Z_m] with value k-1.
        // Note that y=0 corresponds to p(y)=0, which is handled by k=1 case.
        res = count(Z_m - 1, k) + count(N - Z_m, k - 1);
    }
    
    // Store the computed result in the memoization table and return it.
    return memo[state] = res;
}

int main() {
    // Use faster I/O operations.
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    // Precompute the zebra numbers needed.
    precompute_zebra();
   
    int t; // Number of test cases.
    cin >> t;
    while (t--) {
        ll l, r, k; // Input range [l, r] and target zebra value k.
        cin >> l >> r >> k;
        
        // Calculate count for the upper bound r. The memo map stores results.
        ll count_r = count(r, k);
        
        // Calculate count for l-1. Computations might reuse results stored in 'memo' from the count(r, k) call.
        ll count_l_minus_1 = count(l - 1, k);
        
        // The final answer for the range [l, r] is the difference.
        cout << count_r - count_l_minus_1 << "\n";

        // Clear the memoization table after each test case to prevent interference between test cases.
        // This is important for correctness if multiple test cases are processed.
        memo.clear(); 
    }
    return 0;
}

Problem link: https://codeforces.com/problemset/problem/2086/E