Python: 自动选择合适的整数数据类型大小

Question

我正在用Python和numpy分配一个（可能很大的）全是零的矩阵。我打算在里面放入从1到N的无符号整数。

N的值变化很大：可能从1到一百万，甚至更多。

在初始化矩阵之前，我就知道N的值。我该如何选择矩阵的数据类型，以确保它能容纳大小为N的（无符号）整数呢？

而且，我希望选择一个能满足这个要求的最小数据类型。

举个例子，如果N是1000，我会选择np.dtype('uint16')。如果N是240，uint16也可以用，但uint8也可以，它是我能用来存放这些数字的最小数据类型。

这是我初始化数组的方式。我在寻找SOMETHING_DEPENDING_ON_N：

import numpy as np
# N is known by some other calculation.
lbls = np.zeros( (10,20), dtype=np.dtype( SOMETHING_DEPENDING_ON_N ) )

谢谢！

哦，明白了！

刚意识到numpy v1.6.0及以上版本有np.min_scalar_type，文档。哎呀！（虽然这些答案仍然有用，因为我没有1.6.0版本）。

Answer 1

有趣的是，这里是我之前一直在尝试的版本，直到@Ignacio Vazquez-Abrams和@wim发布了他们的答案，使用了位移操作：

def minimal_uint_type(N):
    bases = [8,16,32,64]
    a = [N>>i for i in bases]
    try: dtype = bases[len(np.nonzero(a)[0])]
    except: raise StandardError('{} is really big!'.format(N))
    return dtype

Answer 2

先建立一个最大值和类型的对应关系，然后再找出比N大的最小值。

typemap = {
  256: uint8,
  65536: uint16,
   ...
}

return typemap.get(min((x for x in typemap.iterkeys() if x > N)))

Answer 3

那我们来写一个简单的函数来完成这个任务怎么样？

import numpy as np

def type_chooser(N):
    for dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
        if N <= dtype(-1):
            return dtype
    raise Exception('{} is really big!'.format(N))

下面是一个使用的例子：

>>> type_chooser(255)
<type 'numpy.uint8'>
>>> type_chooser(256)
<type 'numpy.uint16'>
>>> type_chooser(18446744073709551615)
<type 'numpy.uint64'>
>>> type_chooser(18446744073709551616)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "spam.py", line 6, in type_chooser
    raise Exception('{} is really big!'.format(N))
Exception: 18446744073709551616 is really big!