"Digit Queries" CSES Solution giving TLE

Question

I really need a more accurate way to size inf_string. This code works but gets exponentially more inefficient the bigger max(queries) is, and therefore just can't hold up to the gargantuan numbers that CSES throws into it.

Problem: Consider an infinite string that consists of all positive integers in increasing order: 12345678910111213141516171819202122232425... Your task is to process q queries of the form: what is the digit at position k in the string? The first input line has an integer q: the number of queries. After this, there are q lines that describe the queries. Each line has an integer k: a 1-indexed position in the string. For each query, print the corresponding digit.

queries = []
for i in range(int(input())):
    queries.append(int(input()))
# TODO FIND A BETTER WAY TO LIMIT inf_string LENGTH
inf_string = [str(x) for x in range(1, max(queries)+1)]
inf_string = "".join(inf_string)
for k in queries:
    print(inf_string[k-1])

Answer 1

Building or iterating the whole string takes far too long. Instead group the numbers by length:

length 1 digit:  1 to 9 -> 9 numbers => 9 digits total
length 2 digits: 10 to 99 -> 90 numbers => 180 digits total
length 3 digits: 100 to 999 -> 900 numbers => 2700 digits total
length 4 digits: 1000 to 9999 -> 9000 numbers => 36000 digits total
etc

The number of digits per group (9, 180, 2700, 36000, etc) are easy to compute. Skip these groups of equal-length numbers until you reach the right group (the remaining k falls into it). Then just compute the right number inside the group and get the right digit from it.

for _ in range(int(input())):
    k = int(input())

    length = 1
    while k > 9 * 10**(length-1) * length:
        k -= 9 * 10**(length-1) * length
        length += 1
    q, r = divmod(k-1, length)
    print(str(10**(length-1) + q)[r])

Takes about 0.03 seconds of the allowed 1 second.

Or keeping 10**(length-1) in a variable for less duplication:

for _ in range(int(input())):
    k = int(input())

    length = first = 1
    while k > 9 * first * length:
        k -= 9 * first * length
        length += 1
        first *= 10
    q, r = divmod(k-1, length)
    print(str(first + q)[r])

Answer 2

As Peilonrayz pointed out, there is a mathematical approach to the problem.

Let's assume there wasn't.

# TODO FIND A BETTER WAY TO LIMIT inf_string LENGTH
inf_string = [str(x) for x in range(1, max(queries)+1)]
inf_string = "".join(inf_string)

With \$k \le 10^{18}\$, your string can have up to \$10^{18}\$ digits. That's a lot of memory for a program to use for one string.

There is no reason to make the string that large. One can simply generate the string, one value at a time, and return the digits of those strings, yielding the infinite digit string while taking almost no memory.

from typing import Iterable

def infinite_digits() -> Iterable[str]:
    """
    An infinite series of digits from the infinite sequence of positive
    integers. Ie:

    "1" "2" "3" "4" "5" "6" "7" "8" "9" "1" "0" "1" "1" "1" "2" "1" "3" ...
    """

    x = 1
    while True:
        yield from str(x)
        x += 1

To find the \$n^{th}\$ digit, simply generate the sequence, discarding all of the digits until you reach the desired one. There is even a recipe for that in the itertools docs:

from more_itertools import nth

def digit(k: int) -> str:
    return nth(infinite_digits(), k - 1)

With the above two functions, your code can be simplified to:

for _ in range(int(input())):
    k = int(input())
    print(digit(k))

The code will still take a very, very long time, but at least it won't consume much memory.

You can make things more efficient by sorting your queries, pass over the infinite_digit() stream once, and remember the positions of interest. Finally, pass over the original query set again, and output the memorized digit for that query value.

It could be up to 1000 times faster, since you have up to 1000 queries, but it can still take a very long time.

Answer 3

You need to identify a pattern. We should be able to build a fairly simple pattern by grouping numbers in specific ways.

123456789  # One dimensional
         10111213141516171819  # Two dimensional
         20212223242526272829
         30313233343536373839
         ...
         90919293949596979899
                             100101102103104105106107108109  200201202203204205206207208209  # Three dimensional
                             110111112113114115116117118119
                             120121122123124125126127128129
                             ...
                             190191192193194195196197198199
                                                           1000100110021003100410051006100710081009  # Four dimensional

1. We know if \$k\$ is less than 10 the value is \$k\$. To see the pattern look at the one dimensional group and see the index where each number belongs. We can see each digit has only one possible position.

   • \$k \in \{1\} \therefore 1\$
     (if \$k = 1\$ then the digit is 1)
   • \$k \in \{2\} \therefore 2\$
     (if \$k = 2\$ then the digit is 2)

   $$ f_1(k) = k $$

2. To make the pattern easier to see \$k_2 = k - 10\$.
   Now \$f_2(0) = 1\$, \$f_2(1) = 0\$, \$f_2(2) = 1\$, \$f_2(3) = 1\$.

   If we focus on how to build each section of the number we can see similar patterns. For example the end digit is 0 in the 2nd column in the entire two dimensional block. Since we know the width of the two dimensional block is 20 we know every \$20n + 1\$ number results in 0.

   • End digits (Odd numbers, \$k_2\ \%\ 2 = 1\$)

     Just like we did for 0 we can look at every second column and see the value repeats all the way down.

     • \$k_2 \in \{1, 21, 41, ..., 161\} \therefore 0\$
       (if \$k_2 = 1\$ or \$k_2 = 21\$ or etc, the digit is 0)
     • \$k_2 \in \{3, 23, 43, ..., 163\} \therefore 1\$
     • \$k_2 \in \{5, 25, 45, ..., 165\} \therefore 2\$

     We can then calculate the digit from the input number. First we can modulo the number, \$k_2 % 20\$, to truncate each of the sets to the same number. For example \$21 \% 20 = 1\$, \$41 \% 20 = 1\$, etc. Now the sets are much simpler.

     • \$k_2\ \%\ 20 \in \{1\} \therefore 0\$
       (if \$k_2\ \%\ 20 = 1\$ the digit is 0)
     • \$k_2\ \%\ 20 \in \{3\} \therefore 1\$
     • \$k_2\ \%\ 20 \in \{5\} \therefore 2\$

     Now getting the same digit for each input is really easy. We can just floor divide by 2. Since \$\lfloor\frac{1}{2}\rfloor = 0\$, \$\lfloor\frac{3}{2}\rfloor = 1\$, etc.

     So we know if the number is odd we can get the digit by using the following formula:

     $$\lfloor\frac{k_2\ \%\ 20}{2}\rfloor$$

   • Head digits (Even numbers, \$k_2\ \%\ 2 = 0\$)

     Just like we did for the columns we can see a pattern in the rows. Every two columns the first number is the same. We can build the sets for where each number is.

     • \$k_2 \in \{ 0, 2, 4, ..., 18\} \therefore 1\$
     • \$k_2 \in \{20, 22, 24, ..., 38\} \therefore 2\$
     • \$k_2 \in \{40, 42, 44, ..., 58\} \therefore 3\$

     We can see any less than 20 are 1, then any less than 40 are 2. We can floor divide by 20 and add 1. And so the formula is:

     $$\lfloor\frac{k_2}{20}\rfloor + 1$$

   We can merge both the previous functions into one function.

   $$ f_2(k_2) = \begin{cases} \lfloor\frac{k_2}{20}\rfloor + 1, & \text{if $k_2 \% 2 = 0$ (head)} \\ \lfloor\frac{k_2\ \%\ 20}{2}\rfloor, & \text{if $k_2 \% 2 = 1$ (end)} \end{cases} $$

3. To make the pattern easier to see \$k_3 = k_2 - 180\$.

   Again we'll focus on building each digit.

   • End digits \$k_3\ \%\ 3 = 2\$

     Ok the end digits are following the same pattern again. Every \$30n + 2 = 0\$, \$30n + 5 = 1\$.

     • \$k_3 \in \{2, 32, 62, ..., 2672\} \therefore 0\$
     • \$k_3 \in \{5, 35, 65, ..., 2675\} \therefore 1\$
     • \$k_3 \in \{8, 38, 68, ..., 2678\} \therefore 2\$

     Again we can use a similar algorithm as before to truncate the sets to just one number each. Afterwards we can figure out how to convert each of the numbers to the digit.

     We can modulo \$k_3\$ by 30 to truncate each set to one number each, leaving 2, 5, 8, .... The final numbers can be converted to the correct output by just floor dividing by 3.

     $$\lfloor\frac{k_3\ \%\ 30}{3}\rfloor$$

   • Mid digits \$k_3\ \%\ 3 = 1\$

     The mid digits have an interesting pattern. Because the numbers are repeating in the 3rd dimension each time starting from around 100, 200, 300, etc. As such I have used nested sets to show the split patterns. The nested sets are just to show the start, end and step of the repeating patterns.

     • \$k_3 \in \{\{1, 4, 7, ..., 28\}, \{301, 304, ... 328\}, ..., \{2401, 2404, ..., 2428\}\} \therefore 0\$
     • \$k_3 \in \{\{31, 34, 37, ..., 58\}, \{331, 334, ... 358\}, ..., \{2431, 2434, ..., 2458\}\} \therefore 1\$
     • \$k_3 \in \{\{61, 64, 67, ..., 88\}, \{361, 364, ... 388\}, ..., \{2461, 2464, ..., 2488\}\} \therefore 2\$

     Just like in both the previous end digits we can use modulo to reduce the size of the set. We can see the outer sets repeat every 300 digits so using modulo 300 we can reduce the sets to something more manageable.

     • \$k_3 \% 300 \in \{1, 4, 7, ..., 28\} \therefore 0\$
     • \$k_3 \% 300 \in \{31, 34, 37, ..., 58\} \therefore 1\$
     • \$k_3 \% 300 \in \{61, 64, 67, ..., 88\} \therefore 2\$

     Now the sets look like sets we got in the head digit. Now we can just floor divide by 30 to get the digit.

     $$\lfloor\frac{k_3\ \%\ 300}{30}\rfloor$$

   • Head digits \$k_3\ \%\ 3 = 0\$

     • \$k_3 \in \{ 0, 3, 6, ..., 297\} \therefore 1\$
     • \$k_3 \in \{300, 303, 306, ..., 597\} \therefore 2\$
     • \$k_3 \in \{600, 603, 606, ..., 797\} \therefore 3\$

     Again the head digit pattern is the same, just with a different number. We can just floor divide by 300 and add one to get the digit.

     $$\lfloor\frac{k_3}{300}\rfloor + 1$$

   Again we can merge the patterns into a single function:

   $$ f_3(k_3) = \begin{cases} \lfloor\frac{k_3}{300}\rfloor + 1, & \text{if $k_3\ \%\ 3 = 0$ (head)} \\ \lfloor\frac{k_3\ \%\ 300}{30}\rfloor, & \text{if $k_3\ \%\ 3 = 1$ (mid)} \\ \lfloor\frac{k_3\ \%\ 30}{3}\rfloor, & \text{if $k_3\ \%\ 3 = 2$ (end)} \end{cases} $$

We can see a pattern emerging. Lets have \$k_1 = k - 1\$.

• We can see similar values 2, 20, 3, 30 and 300. So we can move the dimension, \$n\$, out of the numbers and get \$1n\$, \$10n\$, \$1n\$, \$10n\$ and \$100n\$ respectively.

• Looking at the previous numbers we can see a pattern again 1, 10 and 100. The pattern between the numbers is \$10^?\$ where we currently don't know what \$?\$ is.

• We can see the numerator's modulo is always \$? + 1\$ the denominator's \$?\$.

As such we can change the above functions to a very similar pattern.

$$ f_3(k_3) = \begin{cases} \lfloor\frac{k_3\ \%\ (3 * 10^{3 - 0})}{3 * 10^{3 - 1}}\rfloor + 1, & \text{if $k_3\ \%\ 3 = 0$} \\ \lfloor\frac{k_3\ \%\ (3 * 10^{3 - 1})}{3 * 10^{3 - 2}}\rfloor, & \text{if $k_3\ \%\ 3 = 1$} \\ \lfloor\frac{k_3\ \%\ (3 * 10^{3 - 2})}{3 * 10^{3 - 3}}\rfloor, & \text{if $k_3\ \%\ 3 = 2$} \end{cases}\\ f_2(k_2) = \begin{cases} \lfloor\frac{k_2\ \%\ (2 * 10^{2 - 0})}{2 * 10^{2 - 1}}\rfloor + 1, & \text{if $k_2 \% 2 = 0$} \\ \lfloor\frac{k_2\ \%\ (2 * 10^{2 - 1})}{2 * 10^{2 - 2}}\rfloor, & \text{if $k_2 \% 2 = 1$} \end{cases}\\ f_1(k_1) = \begin{cases} \lfloor\frac{k_1\ \%\ (1 * 10^{1 - 0})}{1 * 10^{1 - 1}}\rfloor + 1, & \text{if $k_1 \% 1 = 0$} \end{cases} $$

Here we have found what \$?\$ is. Looking at the "if" we can see \$?\$ is based on the modulo of \$k\$. As such we can improve the function further to:

$$ f_n(k_n) = \lfloor\frac{k_n\ \%\ (n10^{n - (k_n\ \%\ n)})}{n10^{n - (k_n\ \%\ n) - 1}}\rfloor + \begin{cases} 1, & \text{if $k_n \% n = 0$}\\ 0, & \text{otherwise} \end{cases} $$

I'll leave getting \$n\$ and \$k_n\$ from \$k\$ as a challenge for the OP.

We can check the equation works with the following code.

def fn(k, n):
    mod = k % n
    return (
        (
            (k % (n * 10 ** (n - mod)))
            // (n * 10 ** (n - mod - 1))
        )
        + (0 if mod else 1)
    )

for row in range(0, 180, 20):
    digits = [
        fn(index, 3)
        for index in range(row, row + 20)
    ]
    print("".join(str(d) for d in digits))

10111213141516171819
20212223242526272829
30313233343536373839
40414243444546474849
50515253545556575859
60616263646566676869
70717273747576777879
80818283848586878889
90919293949596979899

Note: The function doesn't work correctly on the 'last chunk' of numbers, fn(180, 2) == 10 but when the function loops back around we get the correct numbers fn(200, 2) == 1. Because we don't need to generate \$00\ \to\ 09\$ I defiled the function. To get the 'pure' function remove + (0 if mod else 1) and adjust the input accordingly. I.e start at 20 not 0.

Answer 4

This is a competition programming problem so it's worth thinking about the algorithmic complexity of our approaches. Let n=10^18 be the original length of the string. Generating the entire string one digit at a time will take O(n) time, even if we don't store the whole string, so that approach won't work. Instead, the goal is to compute precisely which number of all the concatenated numbers our index lands in and how far into the number it is. (By the way, this sequence is http://oeis.org/A033307)

A simple approach is to group all numbers by how many digits is in each number, like in Kelly Bundy's answer.

• For numbers 1 to 9, there are 9 numbers and each number is 1-digit, so there are 9*1=9 digits in the group.
• For numbers 10 to 90, there are 90 numbers and each number is 2-digit, so there are 90*2 = 180 total digits.
• etc.
• In general, for d-digit numbers, there are g = 9*(10^(d-1)) * d total digits. This is http://oeis.org/A212704.

So we can find out what d is for our given k by incrementing d and subtracting from k until we cannot subtract any more. Each time we subtract on the order of 10^d, so we only need O(log10(n)) steps to reach our group (check index of http://oeis.org/A033713 to see just how many steps we need), so for q=1000 queries, this is very fast.

Then we can compute the digit inside the number we are currently in by using m = 10^(d-1) + q, where q is floor division of k-1 by d (convince yourself why this is true). Within each number, we need to extract the exact digit, which can be done with mod and floor division powers of 10 operations in constant time instead of using strings. There is little performance difference but I like to avoid using strings for numeric problems if I can as a matter of elegance: although this problem is presented as one of an infinite string, it is computed mathematically without any strings at all, showing that this is not really a string manipulation problem but one of careful counting and nothing more.

Here is the modified version of Kelly Bundy's code, using the new python walrus operator to make the code a little cleaner and without using strings, thus very fast and suitable for implementation in C/assembly/any language without nice string functions.

def f(k):
    d = 1
    while k > (g := 9 * 10**(d-1) * d):
        k -= g
        d += 1
    q, r = divmod(k-1, d)  # convert 1-indexing
    m = 10**(d-1) + q
    return (m // (10**(d-1-r))) % 10