NumPy slice assignment

Most sieves in these solutions do one thing over and over: update every $k$-th entry of a table, starting at some offset. That is a C for loop with a stride, and NumPy can write it as one statement, a[start::step] op= value, which runs the whole loop in compiled code. The NumPy page explains why vectorising pays and where it doesn't. This page covers the one construct that turns up in every problem below: what a slice assignment means, which C loop it replaces, and the details that decide whether the replacement is correct.

The idea

A basic slice a[start:stop:step] is a view: no copy, just a new start pointer and a stride into the same buffer. Assigning to it, or applying a compound operator to it, writes through to the original array. So the two lines below do the same work. One runs in the interpreter's bytecode loop, the other runs entirely in C:

for (long long m = start; m <= n; m += step) a[m] OP= v;

a[start::step] OP= v      # stop omitted: "to the end of the array"

Read the slice as the loop header. start is the initialiser, step is the increment, and the array's length is the m <= n guard, because a slice that runs past the end just stops there and never raises. The C ports of these problems keep this loop verbatim, so each NumPy line in the Python has a matching line in the C file.

Every problem below uses one of four variations on that loop header.

Mark: the sieve of Eratosthenes. flags[p*p::p] = False clears the multiples of $p$ from $p^2$ onward. 216, 249, 319 and 351 all build their prime tables this way. The sieve page explains why the start is $p^2$. On the NumPy side, the Python loop runs only over $p \le \sqrt n$, while each statement does about $n/p$ writes. 304 uses the stride to skip even numbers: it clears the evens once with base[4::2], then marks odd multiples with step $2p$, the same halving the C loop gets from j += 2 * p. 291 halves the memory as well: its table holds only the odd numbers, with index $i$ standing for $2i + 1$. The odd multiples of $p$ are $2p$ apart as values but $p$ apart as indices, so the marking line becomes sieve[p*p // 2::p] = False. The change is in the index map, and the loop header has to be translated through it.

None of this is specific to NumPy. Python's own bytearray and list accept the same extended slice on the left, and 967 builds its small prime list that way, with is_prime[p*p::p] = bytearray(len(range(p*p, b + 1, p))). The built-in form is stricter, though. It has no broadcasting, so the right-hand side must be a sequence of exactly the slice's length (hence the len(range(...))), and it has no compound operators, so it can mark but cannot accumulate.

Accumulate: divisor and multiplicative sieves. Swap = for +=, -=, *= or //= and the strided loop gathers arithmetic instead of marking. With tau[i::i] += 1 every multiple of $i$ counts $i$ as a divisor, which is how 420 builds its divisor-count table. In 351 and 512, phi[p::p] -= phi[p::p] // p applies the factor $(1 - 1/p)$ of Euler's totient to every multiple of $p$ in one statement; 512 doesn't even keep a separate prime table, since an entry still equal to its index is one no earlier prime has touched. The Möbius sieves in 319, 464 and 641 take two strides per prime: they flip the sign at step $p$, then zero the entries at step $p^2$. 319 also runs cofactor[q::q] //= p for each prime power $q$, removing the small primes from every index so that what is left over shows which numbers still have a large prime factor. 632 does the reverse inside each segment of a segmented sieve: product[start::p] *= p multiplies up the small primes of each $d$, and omega += product < d then adds one to $\omega$ wherever the small part falls short, because a squarefree $d$ can have at most one prime factor above the sieving bound. A third stride on the same segment, squarefree[-lo % square::square] = False with square = p*p, removes the non-squarefree $d$. Each of these keeps the same m += step header in its C version, and only the loop body changes.

A slice can also patch a sieve's output after the fact, by residue class instead of by prime. 875 needs a weight that is $\varphi(m)\,m^3$ for most $m$ but is zero when $m \equiv 2 \pmod 4$ and sixteen times larger when $4 \mid m$. It runs the totient sieve unchanged and then writes the two exceptions as w[2::4] = 0 and w[4::4] = w[4::4] * 16 % MOD. Each slice is one branch of a case split on $m \bmod 4$, so the code has no per-element if.

The value on the right need not be a scalar. Any array of the slice's length lines up with it entry by entry, so the loop body can depend on which multiple it is visiting. 675 maintains a running product modulo a prime, and the factor each multiple $kp$ contributes depends on $v_p(kp)$ and on the exponent of $p$ in $(kp)!$. It builds that per-multiple factor as an array of length $n/p$ — itself by strided += 1 passes at steps $p, p^2, \dots$ followed by a cumsum, which is Legendre's formula in two lines — and then applies it as mult[p::p] = mult[p::p] * ratio % MOD. Nor is the operator limited to the arithmetic ones Python has syntax for. 407 keeps a running maximum per index: it takes the view best[s*(s+1)::s] (trimmed by a second slice, which is still a view) and hands it to a ufunc as both input and output, np.maximum(view, candidates, out=view). Any binary ufunc with out= set to a slice is a compound assignment; += is just the one with an operator spelling.

Pair: the half-length divisor count. Looping $i$ up to $n$ costs $n$ Python-level statements, and each one has a fixed overhead whatever its length. Past $i \approx n/2$ the slices are only one or two entries long, so that overhead becomes most of the cost. 179 avoids this with the divisor-pair idea. Every divisor $a \le \sqrt m$ comes with a partner $m/a$, so it runs div_count[a*a::a] += 2 for $a \le \sqrt n$ and corrects the square case by one. The work is the same, but there are $\sqrt n$ statements instead of $n$. 485 writes the same sieve without the correction. It adds 1 at the square, d[a*a] += 1, and starts the stride one step later, d[a*(a+1)::a] += 2, so the square is never counted twice. The C version doesn't need this trick, because a loop iteration there costs a few cycles. The trick exists because in Python the cost per slice dominates.

Offset: residue classes and windows. Nothing forces the start to be a multiple of the step. A slice can describe any arithmetic progression, so a congruence turns into a slice once you know its least solution:

  • 216 sieves the values $2n^2 - 1$ rather than the integers. Each prime $p$ divides that polynomial exactly when $n$ is one of the two square roots of $1/2 \pmod p$, so each prime contributes two slices, is_t_prime[r::p] for each root $r$. Finding $r$ is a modular square root. Once $r$ is known, the marking is the same one-line stride. 291 does the same for $2n^2 + 2n + 1$, with one addition. The root class of $q$ contains the single $n$ with $f(n) = q$, which is prime and must survive, so where the least root is that $n$ the start moves one period further on, to $r + q$.
  • 659 needs more than divisibility. It strips the whole power of $p$ from each value of $4k^2 + 1$, so that what is left is the largest prime factor. Each power $q = p^e$ gets its own slice, values[k::q] //= p. The start $k$ is the root modulo $p^e$, which it gets from the root modulo $p^{e-1}$ by a Hensel lift. So the residue class shrinks as the step grows, and the loop stops when the lifted root passes the end of the array.
  • 304 runs a segmented sieve over a window that starts at low. The first multiple of $p$ inside it is at offset (-low) % p, which is the whole adjustment the segment needs.
  • 256 counts, for each area in a window, the rooms of that area that can be tiled. For one side length $a$, the valid areas form a band of an arithmetic progression, and the step doubles to $2a$ when an odd side forces the other side to be even. One counts[lo:hi:step] += 1 records a whole band, so the window is filled with a Python loop over $(a, k)$ and never over the areas themselves.
  • 967 does not sieve the range at all. For each $m$ it needs the sum, modulo 3, of the distinct primes from a fixed small set that divide $m$. That value depends only on $m \bmod Q$, where $Q$ is the product of the set. So it sieves a single period, residue[::p] += p % 3 over an array of length $Q$, takes a cumsum of each class, and answers a count up to any bound $x$ as whole periods $\lfloor x/Q \rfloor$ plus a prefix lookup at $x \bmod Q$. The range is far too long to allocate, but it is periodic, so one period of slices covers it.

Shift: the overlapping slice. The fourth variation has no stride. The source is a slice of the same array, shifted:

counts[p:] += counts[:-p]        # every subset sum s gains the subsets of s - p

for (long long s = total; s >= p; s--) counts[s] += counts[s - p];

This is the 0/1 knapsack step in 249. The C loop runs downwards on purpose: going upwards would read entries that this same prime has already updated, and an item could be used twice. NumPy gets the downward loop's result without any special care, because it detects the overlap and behaves as if the right-hand side were copied before the write. The NumPy page also covers this case. What matters here is the consequence: the slice gives you only the copy semantics. If a recurrence really needs the upward behaviour, where each entry reads its already-updated neighbour (the unbounded knapsack, a running total), it is a sequential dependency, and one slice statement cannot express it. That needs np.cumsum, or a Python loop.

The copy semantics carry over when the source is gathered rather than shifted. 452 and 668 both count primes with Lucy_Hedgehog's method, which keeps $\pi$-style counts at every value $\lfloor n/k \rfloor$ and, for each prime $p$, updates them by

small[p2:r+1] -= small[v // p] - sp      # v = p*p .. r; sp = primes below p

The source index $v/p$ can itself be $\ge p^2$, an entry this same pass updates. The scalar algorithm handles that by walking $v$ downwards so every read sees the pre-update value — the same reason 249's C loop runs downwards — and the NumPy line gets it for free, because the gather small[v // p] is a fresh array, evaluated in full before anything is written. The comment in 452 says exactly that. The large half of the table uses the stride on the source side instead: large[p:split*p+1:p] reads every $p$-th entry and feeds a contiguous target. 521 runs the same two updates on a second pair of tables, which hold prime sums modulo $M$ instead of counts. The slices are the same, and only the right-hand side changes.

A shift can run along one axis of a 2-D table. 294 keeps counts indexed by (residue mod 23, digit sum), and joining two digit blocks convolves the tables in both axes, but differently. Digit sums add and are cut off at 23, so that axis is a truncated shift, out[:, s:] = (out[:, s:] + count * rolled[:, :D+1-s]) % MOD. Residues add modulo 23, so that axis must wrap around, which no slice can express. It uses np.roll instead, which builds a shifted copy. The rule is: a linear shift is a view, a cyclic shift is a copy.

Views can also be reversed and reshaped. 485 needs the maximum of every length-$k$ window. It pads the array to a multiple of $k$ and reshapes it into $k$-wide blocks, which is still a view. It takes running maxima forwards along each block, and also backwards through the reversed view grid[:, ::-1]. A window spans at most two adjacent blocks, the end of one and the start of the next, so its maximum is np.maximum(suffix[:count], prefix[k-1:k-1+count]), two views aligned by an offset of $k - 1$. That is the van Herk/Gil–Werman trick, at $O(1)$ work per window and with no Python loop.

The same alignment gives cheap comparisons between neighbours. In 179, div_count[2:n] == div_count[3:n+1] lines up $d(m)$ with $d(m+1)$ for every $m$ at once. In 464, key[1:] != key[:-1] finds where a sorted key changes value. Both read two views of the same buffer and write nothing back.

How to reason about it

Check the loop header first, then the slice. Most errors here are off-by-one mistakes in start or stop. Write the C loop you mean, including its first index, step and final index, and then translate it. A missing stop means "through index len - 1", so a table that must include $n$ needs length $n + 1$. That is why every array in these solutions has size limit + 1. An explicit stop is exclusive: 256's band slice ends at last - lo + 1 to make it inclusive.

Edge cases do not raise errors. A start past the end gives an empty slice and a silent no-op. 216 checks start > max_n explicitly because otherwise nothing would show that a root was skipped. The shift form has a trap at zero: a[:-0] is a[:0], which is empty, not the whole array. A zero shift then fails with a shape mismatch instead of doubling the array. Shift by a value you know is at least 1.

Compound assignment keeps the left-hand dtype and wraps. a[s::k] += x casts the result back to a's dtype under the same_kind rule. An int8 Möbius table that gets += 300 wraps without any warning. That is safe for $\mu \in \{-1, 0, 1\}$ and for 256's uint16 counts, whose values stay tiny. It is a problem for a knapsack count that doubles every pass, which is why 249 reduces modulo its bound before a uint64 could overflow, and why 675 keeps every operand of mult[p::p] * ratio % MOD below $2^{30}$, so the product fits int64 before the reduction. Reductions are the other way round: sum and cumsum promote small integer types to the platform integer, so 641's Mertens prefix over an int8 Möbius table does not wrap — it spells dtype=np.int64 anyway, which costs nothing and survives a later change of the table's type. The table's dtype is part of the correctness argument; see watch the integer width.

Basic slices update every entry; fancy indices may not. A stride touches each position at most once, so += on it is exact. An integer-array index such as a[idx] += 1 is buffered, so a repeated index in idx is incremented once, not twice. Use np.add.at or np.bincount when indices can repeat. 304 mixes the two kinds safely. For small primes it uses one strided slice per prime. For large primes, which hit the segment at most once, it uses a single fancy assignment seg[offsets] = False. Duplicates in that array cannot matter because the value assigned is constant. 675's large-prime pass, mult[large * k] = mult[large * k] * ratio_k[k] % MOD, is a fancy update, and it is exact for a different reason: for a fixed $k$, distinct primes give distinct indices $pk$, so no index repeats within one statement. 967 is the case where indices do repeat. It sorts over a million divisor terms into far fewer classes, and adds each term to its class with np.add.at(acc, cls, values). Writing acc[cls] += values there would keep only one term per class and give a wrong total, with no error.

The count of statements is the cost model. A slice statement costs a fixed overhead plus a C loop over its entries. Total work is the usual sieve bound, $\sum n/p = O(n \log\log n)$ for primes or $\sum n/i = O(n \log n)$ for divisors. NumPy wins when the number of statements is much smaller than the work, as with $\pi(\sqrt n)$ marking passes or 179's $\sqrt n$ pair passes. When most slices are a handful of entries long, the fixed overhead dominates and the Python can fall behind a plain list loop. In that case either restructure (pair the divisors, sieve only to $\sqrt n$ as 319's Möbius table does) or accept that the problem belongs in C. 675 shows the restructuring in its general form: primes above $\sqrt n$ with only a few multiples each would mean one short slice per prime, so it transposes the loop — the Python loop runs over the multiplier $k$, and each statement touches the $k$-th multiple of every large prime at once. The work is unchanged; the statement count drops from $\pi(n)$ to a small constant. 291 uses the same transposition for its large primes. It groups them into doubling bands $[L, 2L)$, where each root class has at most $N/L + 1$ members, and for each $k$ marks the $k$-th member of every class in the band with one fancy assignment. When the work genuinely fits the strided shape, a slice assignment keeps a sieve readable and close to C speed: one line per loop, matching the C file line for line.

Problems27 / 27

  • ● 0179 — Consecutive Positive Divisors
  • ● 0216 — The Primality of $2n^2 - 1$
  • ● 0249 — Prime Subset Sums
  • ● 0256 — Tatami-Free Rooms
  • ● 0291 — Panaitopol Primes
  • ● 0294 — Sum of Digits - Experience #23
  • ● 0304 — Primonacci
  • ● 0319 — Bounded Sequences
  • ● 0351 — Hexagonal Orchards
  • ● 0407 — Idempotents
  • ● 0420 — $2 \times 2$ Positive Integer Matrix
  • ● 0452 — Long Products
  • ● 0464 — Möbius Function and Intervals
  • ● 0485 — Maximum Number of Divisors
  • ● 0512 — Sums of Totients of Powers
  • ● 0521 — Smallest Prime Factor
  • ● 0632 — Square Prime Factors
  • ● 0641 — A Long Row of Dice
  • ● 0659 — Largest Prime
  • ● 0668 — Square Root Smooth Numbers
  • ● 0675 — $2^{\omega(n)}$
  • ● 0787 — Bézout's Game
  • ● 0875 — Quadruple Congruence
  • ● 0883 — Remarkable Triangles
  • ● 0967 — $B$-Trivisible Numbers
  • ● 1001 — Connections I
  • ● 1003 — Lonely Singles