07 · 把循环改写成数组运算：六个套路

import numpy as np
import timeit

rng = np.random.default_rng(0)
N = 200_000
x = rng.random(N)

def loop_cumsum(x):                      # 循环版：一个累加变量 + 逐个写回
    out = np.empty(len(x))
    s = 0.0
    for i in range(len(x)):
        s += x[i]
        out[i] = s
    return out

print("最大差", np.abs(loop_cumsum(x) - np.cumsum(x)).max())   # ⭐ 先对正确性

t_loop = min(timeit.repeat(lambda: loop_cumsum(x), number=1, repeat=5))
t_arr  = min(timeit.repeat(lambda: np.cumsum(x),   number=1, repeat=5))
print(f"循环 {t_loop*1e3:.2f} ms | np.cumsum {t_arr*1e3:.3f} ms | {t_loop/t_arr:.0f}x")

import numpy as np
import timeit

rng = np.random.default_rng(0)
N, G = 200_000, 1000
g = rng.integers(0, G, N)          # 组号，取值 0..999
x = rng.random(N)

def loop_group(g, x):                    # 循环版：dict 里 +=
    d = {}
    for gi, xi in zip(g, x):
        d[int(gi)] = d.get(int(gi), 0.0) + float(xi)
    return d

def arr_group(g, x):                     # 数组版：一行
    return np.bincount(g, weights=x, minlength=G)

d, a = loop_group(g, x), arr_group(g, x)
print("最大差", max(abs(d[k] - a[k]) for k in d))

t_loop = min(timeit.repeat(lambda: loop_group(g, x), number=1, repeat=5))
t_arr  = min(timeit.repeat(lambda: arr_group(g, x),  number=1, repeat=5))
print(f"循环 {t_loop*1e3:.2f} ms | np.bincount {t_arr*1e3:.3f} ms | {t_loop/t_arr:.0f}x")

import numpy as np
import timeit

rng = np.random.default_rng(0)
N, k = 200_000, 10
x = rng.random(N)

def topk_sort(x, k):                     # 全排序再切
    return np.argsort(-x)[:k]

def topk_part(x, k):                     # 只保证「前 k 个在前面」，不保证它们有序
    idx = np.argpartition(-x, k)[:k]
    return idx[np.argsort(-x[idx])]      # ⭐ 需要有序的话，再对这 k 个排一次

print("集合一致", set(topk_sort(x, k).tolist()) == set(topk_part(x, k).tolist()))

t_sort = min(timeit.repeat(lambda: topk_sort(x, k), number=1, repeat=5))
t_part = min(timeit.repeat(lambda: topk_part(x, k), number=1, repeat=5))
print(f"argsort {t_sort*1e3:.2f} ms | argpartition {t_part*1e3:.3f} ms | {t_sort/t_part:.1f}x")

import numpy as np
import timeit
from numpy.lib.stride_tricks import sliding_window_view

rng = np.random.default_rng(0)
N, w = 200_000, 50
x = rng.random(N)

def loop_window(x, w):                   # 循环版：切一片算一次均值
    return np.array([x[i:i+w].mean() for i in range(len(x) - w + 1)])

def arr_window(x, w):                    # 数组版：一次性铺开成 (N-w+1, w)
    return sliding_window_view(x, w).mean(-1)

print("最大差", np.abs(loop_window(x, w) - arr_window(x, w)).max())

t_loop = min(timeit.repeat(lambda: loop_window(x, w), number=1, repeat=5))
t_arr  = min(timeit.repeat(lambda: arr_window(x, w),  number=1, repeat=5))
print(f"循环 {t_loop*1e3:.2f} ms | sliding_window_view {t_arr*1e3:.3f} ms | {t_loop/t_arr:.0f}x")

v = sliding_window_view(x, w)
print("展开后形状", v.shape, "| 与原数组共享内存", np.shares_memory(v, x))
print("如果真复制要占", round(v.nbytes/1024**2, 1), "MiB；实际 x 只占",
      round(x.nbytes/1024**2, 1), "MiB")

import numpy as np
import timeit
from bisect import bisect_right

rng = np.random.default_rng(0)
N = 200_000
edges = np.array([0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0])   # ⭐ 必须已排序
x = rng.random(N)

def loop_bucket(x, edges):               # 循环版：一个一个 bisect
    e = edges.tolist()
    return np.array([bisect_right(e, xi) for xi in x])

def arr_bucket(x, edges):                # 数组版：一次全查完
    return np.searchsorted(edges, x, side="right")

print("完全一致", np.array_equal(loop_bucket(x, edges), arr_bucket(x, edges)))

t_loop = min(timeit.repeat(lambda: loop_bucket(x, edges), number=1, repeat=5))
t_arr  = min(timeit.repeat(lambda: arr_bucket(x, edges),  number=1, repeat=5))
print(f"循环 {t_loop*1e3:.2f} ms | np.searchsorted {t_arr*1e3:.3f} ms | {t_loop/t_arr:.0f}x")

import numpy as np
import timeit

rng = np.random.default_rng(0)
N = 200_000
y = rng.normal(size=N)                   # 有正有负的序列

def loop_sign(y):                        # 循环版：比较相邻两个
    out = []
    for i in range(len(y) - 1):
        if (y[i] < 0) != (y[i+1] < 0):
            out.append(i)
    return np.array(out)

def arr_sign(y):                         # 数组版：signbit -> diff -> 非零位置
    return np.flatnonzero(np.diff(np.signbit(y)))

print("完全一致", np.array_equal(loop_sign(y), arr_sign(y)), "| 变号次数", len(arr_sign(y)))

t_loop = min(timeit.repeat(lambda: loop_sign(y), number=1, repeat=5))
t_arr  = min(timeit.repeat(lambda: arr_sign(y),  number=1, repeat=5))
print(f"循环 {t_loop*1e3:.2f} ms | flatnonzero+diff {t_arr*1e3:.3f} ms | {t_loop/t_arr:.0f}x")

import numpy as np
import timeit

k = 10
for N in (200_000, 5_000_000):
    x = np.random.default_rng(0).random(N)
    t_sort = min(timeit.repeat(lambda: np.argsort(-x)[:k],         number=1, repeat=5))
    t_part = min(timeit.repeat(lambda: np.argpartition(-x, k)[:k], number=1, repeat=5))
    print(f"N = {N:>9,} : argsort {t_sort*1e3:7.2f} ms | "
          f"argpartition {t_part*1e3:6.2f} ms -> {t_sort/t_part:.1f}x")