Check numerically if numbers in array start with given digits

Viewed 87

I have a numpy array of integers a and an integer x. For each element in a I want to check whether it starts with x (so the elements of a usually have more digits than x, but that's not guaranteed for every element).

I was thinking of converting the integers to strings and then checking it with pandas

import pandas as pd
import numpy as np

a = np.array([4141, 4265, 4285, 4, 41656])
x = 42

pd.Series(a).astype(str).str.startswith(str(x)).values   # .values returns numpy array instead of pd.Series

[False  True  True False False]

This works and is also quite performant, but for educational purpose I was wondering if there is also an elegant solution doing it numerically in numpy only.

4 Answers

You can get the number of digits using the log10, then divide as integer:

# number of digits of x
n = int(np.ceil(np.log10(x+1)))

# number of digits in array
n2 = np.ceil(np.log10(a+1)).astype(int)

# get first n digits of numbers in array
out = a//10**np.clip((n2-n), 0, np.inf) == x

output: array([False, True, True, False, False])

Usually, we seek a tradeoff between elegant and efficient solutions.

The pandas.Series.str.startswith is faster than the numpy.char.startswith for this test case and both are quite elegant but slower than the numerical approach proposed by @mozway.

Here we have the three solutions:

import numpy as np
import pandas as pd

def pandas_solution(arr, num):
  return pd.Series(arr).str.startswith(str(num)).values

def mozway_solution(arr, num):
  n = int(np.ceil(np.log10(arr + 1)))
  n2 = np.ceil(np.log10(arr + 1)).astype(int)
  out = arr // 10 ** np.clip((n2 - n), 0, np.inf) == num
  return out

def numpy_solution(arr, num):
  return np.char.startswith(arr, str(num))

And here is our test:

import random
import time
import matplotlib.pyplot as plt

labels = ["Pandas", "Mozway", "Numpy Char"]
sizes = [1000, 10000, 10000, 1000000]
fig, axs = plt.subplots(1, 4, figsize=(20, 3))

for counter, ax in zip(sizes, axs):
  pandas_time = []
  mozway_time = []
  numpy_time = []

  arr_int = np.array([random.randint(0, 1000) for _ in range(counter)])
  arr_str = np.array(arr_int, dtype="str")
  
  for trial in range(10): 
    start = time.time()
    pandas_solution(arr_str, 42)
    end = time.time()
    pandas_time.append(end - start)

    start = time.time()
    mozway_solution(arr_int, 42)
    end = time.time()
    mozway_time.append(end - start)

    start = time.time()
    numpy_solution(arr_str, 42)
    end = time.time()
    numpy_time.append(end - start)
  ax.set_title("Size of array: " + str(counter))
  ax.boxplot([pandas_time, mozway_time, numpy_time], labels=labels)

With the following results:

enter image description here enter image description here

You can use numpy.char.startswith as an alternative.

import numpy as np

a = np.array([4141, 4265, 4285, 4, 41656])
x = 42

res = np.char.startswith(a.astype(str), str(x))
print(res)
# [False  True  True False False]

Numerically, the number n = xy...z matches the prefix m = uv...w if, after dividing n by a suitable power of 10 to get the same number of digits, n/10^k = m (integer division).

You have the option of determining k by trial and error (linear search or even dichotomic search), but this is not efficient. Another option is to use logarithms to find the number of digits of n and m.

I would not call such a solution elegant.

Related