TopCoder SRM 591 Editorial: PyramidSequences

PyramidSequences

First, fix some $n$ and $m$. We begin by looking at an easier problem: If instead of looking at sequences of the form $1, 2, \ldots, n - 1, n, n - 1, \ldots, 3, 2$, what if all the sequence elements were distinct (e.g. if the sequence was of the form $1, 2, \ldots, 2 n - 2$)? The Chinese remainder theorem guarantees that a number $x$ from the $n$-sequence will be paired with a number $y$ from the $m$-sequence if and only if $x \equiv y \pmod{2g}$, where $g = \gcd( n - 1, m - 1 )$. We can partition the $n$- and $m$-sequences into equivalence classes modulo $2g$; each class contains $(n - 1) / g$ or $(m - 1) / g$ elements, so there are $(n - 1) (m - 1) / g^2$ pairings per class, or $(n - 1) (m - 1) / g^2$, or $2 (n - 1) (m - 1) / g$ unique pairings.

Unfortunately, the pairings are not unique. Each $k$ in the $n$-sequence has two possible indices modulo $2 g$: $\{ k \% (2 g)$ and $(2 n - k) \% (2 g) \}$. Let us first consider the two indices are equal, i.e. $k \equiv (2 n - k) \pmod{2g}$. Simplifying, this gives us $k \equiv n \pmod g$. Since $g$ is a multiple of $n - 1$, then $n \equiv 1 \pmod g$, so our final expression is $k \equiv 1 \pmod g$; i.e. $k \equiv 1 \text{ or } g + 1 \pmod{2g}$. Note that this condition is independent of $n$ or $m$, so it suffices to simply find the number of $k$ equivalent to $1$ or $g + 1$ modulo $2g$ in each sequence and multiply the counts accordingly.

Now consider the case in which $k$ has two distinct indices, $k \% (2 g)$ and $(2 n - k) \% (2 g)$. Again, since $n \equiv 1 \pmod g$, the second index can be written $(2 - k) \% (2 g)$. If $k \% (2 g) \in [2, g]$, then $(2 - k) \% (2 g) \in [g + 2, 2g]$ and vice versa. Thus, each $k$ appears exactly once in the range $[2, g]$ and exactly once in $[g + 2, 2g]$. It suffices, then, to compute the result for one range. This is simply $(g - 1) (n - 1) (m - 1) / g^2$.

Putting the two cases ($k \equiv 1 \pmod g$ or not) together, we have the following Python code. Some care must be taken to account for the special cases $1$, $n$, and $m$.

from fractions import *

class PyramidSequences(object):
 def distinctPairs(self, n, m):
  g = gcd(n - 1, m - 1)
  an, am = [(x - 1) / g for x in n, m]
  bn, bm = [x / 2 + 1 for x in [an, am]]
  cn, cm = [x / 2 + (x % 2) for x in [an, am]]
  return (g - 1) * an * am + bn * bm + cn * cm