python - Python 中的 strpbrk()

Question

在我编写的一些 Python 代码中，我需要计算字符串中任何一组字符的出现次数。换句话说，我需要计算字符 [c1, c2, c3,...,cn] 在字符串中出现的总次数。

在 C 中，称为 strpbrk() 的函数可用于执行此操作，通常在 x86 处理器上带有特殊指令以使其更快。

在 Python 中，我编写了以下代码，但它是我的应用程序中最慢的部分。

haystack = <query string>
gc_characters = 0
for c in ['c', 'C', 'g', 'G']:
    gc_characters += haystack.count(c)

有没有更快的方法来做到这一点？

Answer 1

我在这里可能有点过火了，但 tl;dr 是原始代码实际上比(编辑:macOS 的)strpbrk() 快，但一些 strpbrk( ) 实现可能会更快!

str.count() 使用 this bundle of strange and beautiful magic在它的胆量中——难怪它很快。

完整代码可在 https://github.com/akx/so55822235 获得

Python代码

这些方法是纯 Python 的，包括 OP 的原始

def gc_characters_original(haystack):
    gc_characters = 0
    for c in ("c", "C", "g", "G"):
        gc_characters += haystack.count(c)
    return gc_characters


def gc_characters_counter(haystack):
    counter = Counter(haystack)
    return sum(counter.get(c, 0) for c in ["c", "C", "g", "G"])


def gc_characters_manual(haystack):
    gc_characters = 0
    for x in haystack:
        if x in ("c", "C", "g", "G"):
            gc_characters += 1
    return gc_characters


def gc_characters_iters(haystack):
    gc_characters = haystack.count("c") + haystack.count("C") + haystack.count("g") + haystack.count("G")
    return gc_characters

Cython 扩展包装 strpbrk()

from libc.string cimport strpbrk

cdef int _count(char* s, char *key):
    assert s is not NULL, "byte string value is NULL"
    cdef int n = 0
    cdef char* pch = strpbrk (s, key)
    while pch is not NULL:
        n += 1
        pch = strpbrk (pch + 1, key)
    return n

def count(s, key):
    return _count(s, key)

...

def gc_characters_cython(haystack_bytes):
    return charcount_cython.count(haystack_bytes, b"cCgG")

手工制作的 C 扩展包装 strpbrk()

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <string.h>

static unsigned int count(const char *str, const char *key) {
  unsigned int n = 0;
  char *pch = strpbrk(str, key);
  while (pch != NULL) {
    n++;
    pch = strpbrk(pch + 1, key);
  }
  return n;
}

static PyObject *charcount_count(PyObject *self, PyObject *args) {
  const char *str, *key;
  Py_ssize_t strl, keyl;

  if (!PyArg_ParseTuple(args, "s#s#", &str, &strl, &key, &keyl)) {
    PyErr_SetString(PyExc_RuntimeError, "invalid arguments");
    return NULL;
  }
  int n = count(str, key);
  return PyLong_FromLong(n);
}

static PyMethodDef CharCountMethods[] = {
    {"count", charcount_count, METH_VARARGS,
     "Count the total occurrences of any s2 characters in s1"},
    {NULL, NULL, 0, NULL},
};

static struct PyModuleDef spammodule = {PyModuleDef_HEAD_INIT, "charcount",
                                        NULL, -1, CharCountMethods};

PyMODINIT_FUNC PyInit_charcount(void) { return PyModule_Create(&spammodule); }

...

def gc_characters_cext_b(haystack_bytes):
    return charcount.count(haystack_bytes, b"cCgG")


def gc_characters_cext_u(haystack):
    return charcount.count(haystack, "cCgG")

---

测量

在我的 Mac 上，计算一百万个随机字母字符串中的 cCgG，即

haystack = "".join(random.choice(string.ascii_letters) for x in range(1_000_000))
haystack_bytes = haystack.encode()
print("original", timeit.timeit(lambda: gc_characters_original(haystack), number=100))
print("unrolled", timeit.timeit(lambda: gc_characters_iters(haystack), number=100))
print("cython", timeit.timeit(lambda: gc_characters_cython(haystack_bytes), number=100))
print("c extension, bytes", timeit.timeit(lambda: gc_characters_cext_b(haystack_bytes), number=100))
print("c extension, unicode", timeit.timeit(lambda: gc_characters_cext_u(haystack), number=100))
print("manual loop", timeit.timeit(lambda: gc_characters_manual(haystack), number=100))
print("counter", timeit.timeit(lambda: gc_characters_counter(haystack), number=100))

产生这些结果:

original               0.34033612700000004
unrolled               0.33661798900000006
cython                 0.6542106270000001
c extension, bytes     0.46668797900000003
c extension, unicode   0.4761082090000004
manual loop           11.625232557
counter                7.0389275090000005

所以，除非我 mac 的 libc 中的 strpbrk() 实现非常低效(编辑:确实如此)，否则最好使用 .count ()。

编辑

我添加了 glibc's strcspn()/strpbrk()而且它比 the näive version of strpbrk() shipped with macOS 快得惊人:

original                       0.329256
unrolled                       0.333872
cython                         0.433299
c extension, bytes             0.432552
c extension, unicode           0.437332
c extension glibc, bytes       0.169704 <-- new
c extension glibc, unicode     0.158153 <-- new

glibc 也有函数的 SSE2 和 SSE4 版本，可能会更快。

编辑 2

自从我顿悟 glibc 的 strcspn() 的智能查找表如何用于字符计数后，我又回到了这里:

size_t fastcharcount(const char *str, const char *haystack) {
  size_t count = 0;

  // Prepare lookup table.
  // It will contain 1 for all characters in the haystack.
  unsigned char table[256] = {0};
  unsigned char *ts = (unsigned char *)haystack;
  while(*ts) table[*ts++] = 1;

  unsigned char *s = (unsigned char *)str;
  #define CHECK_CHAR(i) { if(!s[i]) break; count += table[s[i]]; }
  for(;;) {
    CHECK_CHAR(0);
    CHECK_CHAR(1);
    CHECK_CHAR(2);
    CHECK_CHAR(3);
    s += 4;
  }
  #undef CHECK_CHAR
  return count;
}

结果非常令人印象深刻，比 glibc 实现高出 4 倍，比原始 Python 实现高出 8.5 倍。

original                       | 6.463880 sec / 2000 iter | 309 iter/s
unrolled                       | 6.378582 sec / 2000 iter | 313 iter/s
cython libc                    | 8.443358 sec / 2000 iter | 236 iter/s
cython glibc                   | 2.936697 sec / 2000 iter | 681 iter/s
cython fast                    | 0.766082 sec / 2000 iter | 2610 iter/s
c extension, bytes             | 8.373438 sec / 2000 iter | 238 iter/s
c extension, unicode           | 8.394805 sec / 2000 iter | 238 iter/s
c extension glib, bytes        | 2.988184 sec / 2000 iter | 669 iter/s
c extension glib, unicode      | 2.992429 sec / 2000 iter | 668 iter/s
c extension fast, bytes        | 0.754072 sec / 2000 iter | 2652 iter/s
c extension fast, unicode      | 0.762074 sec / 2000 iter | 2624 iter/s