\section{Introduction}
Many real world applications need to identify events that
repeat very often. Examples of such applications are road traffic peaks~\cite{fha:conj}, load peaks on web servers~\cite{web:02}, monitoring events in computer networks~\cite{bhivw:07}, life event histories~\cite{rbs:13} and many others. Finding such events often leads to useful insights by shedding light on the structure of the data, and giving a basis to predicting future events and behavior. Moreover, in some applications frequent events can point out a problem. In a computer network, for example, repeating error messages can indicate a misconfiguration, or even a security intrusion such as a port scan \cite{MH:ICDE01}.

In Stringology, the problem of counting the occurrences of every subword of length $m$ that appears in text $S$ is a well-known exercise in the power of suffix trees~\cite{W-73} or suffix arrays~\cite{MM-90,KS-03}. However, in reality one seldom finds {\em exact} repetitions of a substring. The situation becomes more complex when we seek the most frequent subword that {\em approximately} occurs in the string.

Let $S[1\ldots n]$ be a text and $k$ and $m$ two integers. The {\em $k$-mappability problem} is defined as follows:
\begin{definition}
For every index $i \in [1\ldots n - m + 1]$, report the number of indices $j$ such that $HD(S[i\ldots i + m - 1], S[j \ldots j + m - 1]) \le k$. With $HD(X,Y)$ denoting the Hamming distance between $X$ and $Y$.
\end{definition}

The $k$-mappability problem was lately tackled by Alzamel et al. ~\cite{DBLP:conf/spire/AlzamelCIKPRS18}. For a text with constant alphabet and $k \in O(1)$, they present an algorithm with linear space and $O(n\log^{k+1}n)$ time. Additionally, they present a quadratic algorithm for reporting the $k$-mappability for a fixed value of $k$ and every $m\in [k\ldots n]$ or a fixed value of $m$ and every $k\in [0\ldots m]$. Finally, they show that the $k$-mappability problem can not be solved in truly subquadratic time unless the Strong Exponentional Time Hypothesis is false. For the case in which $k = 1$ and a constant size alphabet, a faster algorithm with linear space and $O(n\log(n)\log\log(n))$ time was presented in ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}. ~\cite{DBLP:journals/tcs/AlzamelCIPRS20} also presented an algorithm with average case linear time for $k=1$, and provided some experimental results.

\textbf{Our results:}
\begin{enumerate}
    \item By enhancing the techniques of ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}, we construct an algorithm for $k$-mappability with linear space and $O(n\log n)$ time for $k=1$ and infinite integer alphabet. This is an improvement over the $O(n\log^{k+1}n)$ time achieved by ~\cite{DBLP:conf/spire/AlzamelCIKPRS18} for $k \in O(1)$. It also improves the faster $O(n\log(n)\log\log(n))$ time for $k=1$ achieved by ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}. In the settings in which infinite integer alphabet is allowed, our algorithm is optimal.
    \item We present a linear time algorithm for $k$-mappability in the case in which $k=1$, the alphabet size is constant and $m\in \Omega (\sqrt{n})$.

\end{enumerate}

The paper is organized as follows. In Section~\ref{s:pre} we define the basic notions used. Section~\ref{s:constk} presents a linear space $O(n\log n)$ time algorithm for $1$-mappability. In Section ~\ref{s:lspace} we present a linear algorithm for $1$-mappability with constant sized alphabet and $m\in \Omega(\sqrt{n})$.

\section{Preliminaries}\label{s:pre}

Let $\Sigma$ be an alphabet. A {\em string} $S$ over $\Sigma$ is a finite
sequence of characters from $\Sigma$.
By $S[i]$, for $1\leq i\leq |S|$, we denote the $i^{th}$
character of $S$.
The {\em empty string} is denoted by $\epsilon$.
By $S[i \ldots j]$ we denote the string $S[i]\ldots S[j]$ called a {\em
  substring}, or {\em factor}, of $S$ (if $i > j$, then the substring is the empty string).
A substring is called a {\em prefix} if $i=1$ and a {\em suffix} if $j=|S|$.
The prefix of length $j$ is denoted by $S[\ldots j]$,
while by $S[i\ldots ]$ we denote the suffix which starts from index $i$ in $S$.
We say that a string $S$ of length $n$ has a \textit{period} $p$, for some $1 \le p \le \frac{n}{2}$ if $S[i] = S[i+p]$ for every $i \in [1 \ldots n-p]$. {\em The period} of $S$, denoted as $per(S)$, is the smallest $p$ that is a period of $S$.
We say that a substring of $S$, denoted as $A = S[a \ldots b]$ is a \textit{run} with period $p$ if its period is $p$, but $S[a-1] \neq S[a - 1 + p]$ and $S[b+1] \neq S[b+1 - p]$. This means that no substring containing $A$ has a period $p$.
The {\em Hamming distance} of two $n$-length strings, $S_1$, and $S_2$, denoted as $HD(S_1,S_2)$, is the number of indices in which they differ. We say that an $m$-length word $w$ has a $k$-ham occurrence in location $i$ of string $S$ if $HD(w,S[i \ldots i + m - 1]) \le k$.

The longest common prefix (suffix) of two indexes $i,j \in [1\ldots n]$ is the maximal integer $\ell$ such that $S[i \ldots i +\ell - 1] = S[j \ldots j + \ell -1]$ ($S[i - \ell + 1 \ldots i] = S[j - \ell + 1 \ldots j]$). We denote the $LCP(i,j)=\ell$ ($LCS(i,j)=\ell$). LCP and LCS are collectively referred to as longest common extensions (LCE).

The {\em suffix tree}~\cite{W-73} is a useful string data structure.
\begin{definition}\label{d:trie}
Let $S_1,\ldots,S_k$ be strings over alphabet $\Sigma$ and let
$\$\not\in \Sigma$.

An {\em uncompacted trie of strings} $S_1,\ldots,S_k$ is an edge-labeled
tree with $k$ leaves. Every path from the root to a leaf corresponds
to a string $S_i$ with a $\$$ symbol appended to its end. The edges on this path are labeled by the symbols of
$S_i$. Strings with a common prefix start at the root and follow the
same path of the prefix, the paths split where the strings differ.

A {\em compacted trie} is the uncompacted
trie with every chain of edges connected by degree-2 nodes
contracted to a single edge whose label is the concatenation of the
symbols on the edges of the chain.

Let $S=S[1],\ldots,S[n]$ be a string over alphabet $\Sigma$. Let
$\{S_1,\ldots,S_n\}$ be the set of suffixes of $S$, where
$S_i=S[i\ldots ],\ \ i=1,\ldots,n$. A {\em suffix tree}
of $S$ is the compacted trie of the suffixes $S_1,\ldots,S_n$.
\end{definition}

For every node $u$, we call the concatenation of the labels on the path from the root to $u$ the \textit{locus} of $u$ denoted as $\mathcal{L}(u)$. For an edge $e$ in the compact trie, we use the same notation $\mathcal{L}(e)$ to denote the label (or the locus) of $e$.  Finally, for a downwards path $P$ in the compact trie, the locus $\mathcal{L}(P)$ is the concatenation of the loci of the edges in $P$. In a compact trie, an edge $e$ can have label s.t. $|\mathcal{L}(e)| > 1$. We refer to the symbol $\mathcal{L}(e)[1]$ as the symbol of $e$.

\begin{theorem}\label{t:W-3} {\em [Weiner~\cite{W-73}]}
For finite alphabet $\Sigma$, the suffix tree of a length-$n$ string
can be constructed in time $O(n)$. For general alphabets it can be
constructed in time $O(n\log \sigma)$, where $\sigma = {\rm
  min}(|\Sigma|, n)$.
\end{theorem}

The suffix tree can be preprocessed in $O(n)$ \cite{bf:04} to be used as a data structure for $LCE$ queries with $O(1)$ query time.

We assume that every node $u$ in the suffix tree contains some auxiliary information about $\mathcal{L}(u)$, that is the number of occurrences of $\mathcal{L}(u)$ in the text $S$ and a pointer to the list of indices in which $\mathcal{L}(u)$ occurs. This information can be evaluated for all the nodes of a given suffix tree $ST$ in $O(|ST|)$ time and require an additional $O(|ST|)$ space.

Over finite alphabet $\Sigma$, the adjacency list of a node $u\in ST$ is represented as an array $A_u[1\ldots |\Sigma|]$ with the edge with symbol $\sigma \in \Sigma$ in $A_u[\sigma]$ (or an emptiness indicator if there is no edge with that symbol).

Over infinite alphabet, the adjacency list of $u\in ST$ is represented as a balanced search tree storing the edges emerging from $u$ in a sorted order of their symbols. In our algorithm, we assume that the representation of the adjacency list allows linear time DFS iteration on the subtree rooted in a node $u \in ST$. This is indeed the case for most balanced trees.

\begin{definition}\label{d:sarray}
The {\em suffix array} of a string $S$, denoted as $SA(S)$, is an integer
array of size $n+1$ storing the starting positions of all
(lexicographically) sorted non-empty suffixes of $S$, i.e.~for all
$1 < r \leq n+1$ we have $S[SA(S)[r-1] .. n] < S[SA(S)[r] ..
  n]$. Note that the empty suffix is explicitly added to the array.
\end{definition}

The suffix array of $S$ corresponds to a pre-order traversal of all the
leaves of the suffix tree of $S$. Various algorithms exist for efficient time and space construction of the suffix array\cite{s:dcc:98,KS-03,FGGV06}. In particular, the suffix array over a fixed finite alphabet can be constructed in linear time.

\section{$O(n\log(n))$ time and $O(n)$ space algorithm for $k = 1$ }\label{s:constk}

\subsection{An overview of the $O(n\log(n)\log\log(n))$ algorithm for $k=1$}
We start with an overview of the ideas for the $O(n\log(n)\log\log(n))$ algorithm of  ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}. They present an algorithm for counting the number of occurrences with \textbf{exactly} one mismatch, for every word of size $m$. Since there is a textbook algorithm for counting the number of exact occurrences of every word, this is sufficient for solving the $1$-mappability problem.

They start by evaluating the suffix tree $T$ of $S$ and trimming the tree at word length $m$. That is, every node $v$ with $|\mathcal{L}(v)| > m$ is removed. Implicit nodes with $|\mathcal{L}|=m$ are made explicit leaves in the trimming process. They proceed to evaluate the heavy paths decomposition of $T$.
\begin{definition}[Heavy Path Decomposition]
Let $T$ be a rooted tree. For every non-leaf vertex $u$, the edge $(u,v)$ is \textit{heavy} if $|I_u| < 2|I_v|$ with $I_x$ denoting the set of leaves in the subtree rooted in the vertex $x$. An edge that is not heavy is called a \textit{light} edge. The heavy path of a vertex $v$ is the maximal path of heavy edges going through $v$ (it may contain 0 edges). For every heavy path $P$, a vertex $u \in P$, and a light edge $(u,v)$ emerging from $u$, we call $T(v)$ a \textit{sidetree} of $P$ (emerging from $u$).
\end{definition}

It is easy to observe that every root-to-leaf path in $T$ consists of at most $\log(n)$ heavy paths and $\log(n)$ light edges. The following observation is the key for the complexity achieved by ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}:

\begin{observation}\label{obs:ham1split}
For every $w= S[i\ldots i + m - 1]$, every 1-ham occurrence of $w$ $w'=w[1 \ldots x-1] \sigma w[x+1 \ldots m]$ with a mismatch in index $x$ corresponds to a node $u$ in $T$ with $\mathcal{L}(u) = w[1\ldots x-1]$. $u$ must have two edges $e_1,e_2$ s.t. there is a downwards path starting with $e_1$ (resp. $e_2$) and ending in a leaf with path label $w[x \ldots m]$ (resp. $\sigma w[x+1 \ldots n]$).
\end{observation}

Consider the following procedure:
For every node $u\in T$ with path label $w$, let the heavy edge emerging from $u$ be $e_h$ with label $d$. Inspect every light edge $e=(u,v)$ with label $c$ emerging from $u$. For every leaf $z\in T(v)$ with label $L(z) = w \cdot c \cdot w_z$ and for every $c' \neq c \in \Sigma$, find the leaf $z'$ with label $\mathcal{L}(z') = w \cdot c' \cdot w_z$, if it exists. If it does, add the number of occurrences of $\mathcal{L}(z')$ to a counter associated with $z$. For the leaf $z_d$ with $\mathcal{L}(z_d) = w \cdot d \cdot w_z$, also increment a counter associated with $z_d$ by the amount of occurrences of $\mathcal{L}(z)$.

It is straightforward from Observation \ref{obs:ham1split} that for every index $i$, every $1$-ham occurrence is counted by the above procedure. As for complexity - every leaf $z$ is iterated once per light edge in the path from the root to $z$. A single iteration on a leaf $z$ consists of a constant number of counter increments and a single query for finding $z'$ with $L(z') = w \cdot c' \cdot w_z$ per symbol $c' \in \Sigma$. Since $|\Sigma| = O(1)$, the bottleneck of the iteration is finding $z'$. The following is proven in ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}:

\begin{theorem}\label{t:conque}
A text $S[1\ldots n]$ can be preprocessed in time $O(n\log\log n)$ and linear space to allow the following query in $O(\log\log n)$ time:

Given a node $u$ in the suffix tree of $S$ with $L(u) = w_1 \cdot c \cdot w_2$ ($w_1,w_2 \in \Sigma^*$ and $c \in \Sigma$) and a symbol $c \neq c' \in \Sigma$, find the node $u'$ with $L(u') = w_1 \cdot c' \cdot w_2$ if it exists.
\end{theorem}

We call the queries described in Theorem \ref{t:conque} \textit{concatenation} queries.

With Theorem \ref{t:conque} the final complexity is clear - every leaf is iterated $O(\log n)$ times and the iteration costs $O(\log\log n)$ after an $O(n\log\log n)$ preprocessing time. The overall time complexity is $O(n\log(n)\log\log(n))$

\subsection{Linear space $O(n\log n)$ algorithm for $k=1$}\label{s:1km}
\textbf{Intuition:} Our algorithm is based on the ideas of ~\cite{DBLP:journals/tcs/AlzamelCIPRS20}. For every light edge $(u,v)$ we iterate every leaf $z\in T(v)$ and wish to find the vertices corresponding to a 1-ham occurrence of $\mathcal{L}(z)$ with a mismatch in index $|\mathcal{L}(u)| + 1$. Instead of using concatenation queries, we construct a lexicographically sorted array of the words $W$ we need to find. Given the sorted array of words, finding the vertices corresponding to these words in the suffix tree can be done in $O(|W|)$. If we manage to construct this sorted array in $O(|W|)$, the amortized time for inspecting a leaf is constant (rather than $O(\log\log n)$).

\textbf{Terminology}. Let $P=(u_1,u_2, \ldots u_x)$ be a heavy path in the heavy path decomposition of the suffix tree $ST$ of $S$. Let $\mathcal{L}(u_i) = w_i$ and let $e_i = (u_i, u_{i+1})$ be the $i$'th heavy edge in $P$ with symbol $d_i$. Let $(u_i,v)$ be a light edge emerging from $u_i$ with symbol $c$ and let $z\in T(v)$ be a leaf. It holds that $\mathcal{L}(z) = w_i \cdot c \cdot s_z$ for some suffix $s_z \in \Sigma^*$.

\begin{definition}
The node $z'\in ST$ is a \textit{$P$-light occurrence} of $z$ if $L(z') = w_i \cdot c' \cdot s_z$ for some $c' \in \Sigma \setminus \{c,d_i \}$. We call the word $hw(z) = w_i \cdot d_i \cdot s_z$ the \textit{$P$-heavy word} of $z$. The node $z'\in ST$ is a \textit{$P$-heavy occurrence} of $z$ if $L(z') = hw(z)$.
\end{definition}

Note that the above definitions are with respect to a heavy path $P$. $z$ may be a leaf in the sidetrees of multiple heavy paths. In every such path, the $P$-light occurrences, $P$-heavy occurrence and the $P$-heavy word of $z$ are different. Also note that the $P$-heavy word and the $P$-heavy occurrence are undefined for leaves in the sidetrees emerging from $u_x$, as the last heavy edge in $P$ is $e_{x-1}$.

In our algorithm, we count the $P$-heavy occurrences and the $P$-light occurrences of every node $z$ in a sidetree of $P$ independently. For every heavy occurrence $z'$, we also count the occurrences of $w(z)$ as $1$-ham occurrences of $w(z')$. We do this for every heavy path $P$. Surely, this process counts all the $1$-ham occurrences.

We start by showing how to efficiently count the $P$-light occurrences.

\begin{observation}\label{obs:lighoc}
For every vertex $z$ with $\mathcal{L}(z) = w$ in a sidetree $T(v)$ emerging from $u_i$, all the $P$-light occurrences of $z$ are also leaves in (different) sidetrees emerging from $u_i$. Furthermore, a leaf $z'$ with $\mathcal{L}(z') = w'$ in a sidetree $T(v') \neq T(v)$ emerging from $u_i$ is a light occurrence of $z$ iff $w[|w_i + 2| \ldots m] = w'[|w_i+2|\ldots m]$.
\end{observation}

Observation \ref{obs:lighoc} is directly derived from the definition of a light occurrence. For every $u_i\in P$, we wish to construct a sorted array consisting of the suffixes starting in index $|w_i| + 2$ of the labels of the leaves of the sidetrees emerging from $u_i$.

We present the following routine:

{\bf Suffix Sorting:}\\
\fbox{\begin{minipage}{12cm}
 As a preprocess procedure, construct the suffix array $SA$ of $S$.
\\\\
{\bf Initialize} an array $A$ of size $n$ consisting of empty lists
\\\\
\textbf{ Alignment step:} Iterate the leaves in the sidetrees of $u_i$. For every leaf $z$, extract $j_z$ - a starting index of $L(z)$. We add $z$ to $A[j_z + |w_i| + 2]$.
\\\\
{\bf Insertion step:} Initialize an empty list $L$. Iterate $SA$ from left to right. When iterating $SA[j]$, add all the nodes in $A[SA[j]]$ to the end of $L$.
\end{minipage}}

\begin{claim}\label{c:sorLight}
After running Suffix Sorting, $L$ is sorted by the lexicographic order of the suffixes starting in index $|w_i| +2$ of $z$. The running time is $O(n + |SE|)$ with $SE$ being the set of sorted elements.
\end{claim}
\textbf{Proof:} leaf $z$ with $\mathcal{L}(z)=w_z$ occurring in index $j_z$ is inserted to $L$ before the leaf $y$ with $L(y) = w_y$ occurring in index $j_y$ only if the suffix of $S$ starting in $j_z + |w_i| + 2$ is lexicographically smaller or equal to the suffix of $S$ starting in $j_y + |w_i| + 2$. Therefore, it can not be the case that $w_z[|w_i| + 2 \ldots m] >_L w_y[|w_i| + 2 \ldots m]$.
\\\\
As for complexity- the alignment step takes $O(|SE|)$ time as it executes a constant amount of list insertions and basic arithmetic operations for every leaf. The insertion step takes $O(|SE| + n)$ time as it iterates over the entire suffix array. The sum of the sizes of the lists in $A$ is identical to the amount of iterated leaves in the alignment step. We assume that the suffix array was evaluated prior to the run of Suffix Sorting. Therefore, we exclude the complexity of computing the suffix array from our running time.   \hfill\ensuremath{\Box}

This is not exactly what we want. If we  execute Suffix Sorting for every node, the $n$ factor will dominate the complexity and the overall time will be quadratic. To avoid that, we present the following algorithm for sorting a batch of sidetrees.

{\bf Batched Suffix Sorting:}\\
\fbox{\begin{minipage}{12cm}
{\sf As a preprocess procedure, construct the suffix array $SA$ of $S$. }.
\\\\
{\sf Input:} A batch of vertices $v_1,v_2 \ldots v_b$
\\\\
{\sf Initialize} an array $A$ of size $n$ consisting of empty lists.
\\\\
{\sf Batched Alignment step:}
\\\\
{\sf For every} $i \in [1\ldots b]$:
\begin{enumerate}
    \item {\sf Initialize} an empty list $L_i$
    \item Iterate the leaves in the sidetrees of $v_i$ with $L(v_i) = w_i$. For every leaf $z$, extract $j_z$ - a starting index of $\mathcal{L}(z)$. Add the pair $(z,L_i)$ to $A[j_z + |w_i| + 2]$.
\end{enumerate}
{\bf Batched Insertion step:}  Iterate $SA$ from left to right. When iterating $SA[j]$, for every  $(z,L) \in A[SA[j]]$, add $z$ to the end of $L$.
\end{minipage}}

The same arguments as in the proof of claim \ref{c:sorLight} can be made to prove the following:
\begin{claim}
After running Batched Suffix Sort, every list $L_i$ has the leaves in the sidetrees of $v_i$ sorted by the lexicographic order of the suffixes starting in index $|w_i| +2$ of $z$. The running time is $O(n +|SE|)$ with $SE$ being the set of sorted elements in the batch.
\end{claim}

To sort the sidetrees in amortized linear time, we set a counter $se=0$ for the amount of leaves in the sidetrees that need to be sorted and an empty list $Sort$. We iterate the vertices in $ST$. For every vertex $u$, we count the number of leaves in the sidetrees of $u$, add this number to $se$ and add $u$ to $Sort$. Once $se > n$, we execute Batched Suffix Sort on $Sort$.

Since the number of leaves in the sidetrees of a vertex $u$ never exceeds $n$, it is guaranteed that $se \le 2n$ when we execute Batched Suffix Sorting. Therefore, the overall complexity is $O(n + 2n) = O(n)$. Since we only execute the batched insertion with $se \ge n$, the amortized time for placing every leaf in the sorted list is constant.

Once we have the sorted list $L=z_1,z_2 \ldots z_t$ of the leaves in the sidetrees emerging from $u_i \in P$, a simple iteration can be implemented to count the number of $P$-light occurrences for every node in $L$. We start by preprocessing $S$ for constant time lcp queries. We iterate $L$. For every consecutive pair of leaves $z_a$ and $z_{a+1}$ with $\mathcal{L}(z_a) = s_1$ occurring in $j_1$ and $\mathcal{L}(z_{a+1}) = s_2$ occurring in $j_2$, we query $l = lcp(j_1 + |w_i| + 2,j_2 + |w_i| + 2)$. If we have $l \ge  m - |w_i| - 1$, then $z_{a}$ and $z_{a+1}$ are $P$-light occurrences of each other (Observation \ref{obs:lighoc}). Once we identify a pair $z_a,z_{a+1}$ of $P$-light occurrences, we proceed in $L$ until we reach a leaf $z_{b+1}$ that is not a $P$-light occurrence. Of course, all the pairs $z_x,z_y$ with $x \neq y$ and $x,y \in [a \ldots b]$ are $P$-light occurrences of each other. We evaluate the sum $Oc$ of occurrences of $\mathcal{L}(z_x)$ for $x \in [a\ldots b]$ and increment the counter of ham-1 occurrences of $z_x$ by $Oc - Oc(z_x)$ with $Oc(z_x)$ being the number of occurrences of $\mathcal{L}(z_x)$.

It can be easily verified that the iteration is linear. For every leaf we execute a single lcp query and a constant number of basic arithmetic operations. We conclude the handling of $P$-light occurrences with the following theorem:

\begin{theorem}\label{t:fastlight}
The $1$-ham occurrences of $\mathcal{L}(z)$ that are corresponding to $P$-light occurrences of some heavy path $P$ can be computed for every leaf $z\in ST$, in $O(n\log n)$ time and linear space.
\end{theorem}

\textbf{Proof:}
For every heavy path $P$ and vertex $u_i \in P$, we compute the sorted list of the suffixes starting in $|w_i| + 2$ of the words of the leaves of the sidetrees emerging from $u_i$. We use the sorted list to find the $P$-light occurrences of the leaves in the sidetrees of $u_i$. Sorting the leaves is done using Batched Suffix Sorting with batches of size between $n$ and $2n$ and takes a constant amortized time per sorted leaf. There may be one 'remainder' batch with size $se < n$ that takes an additional $O(n)$ time to sort. Given the sorted lists, finding the $P$-light occurrences is linear in the number of leaves in the sidetrees of $u_i$. Every leaf $z$ participates in at most $\log(n)$ different sidetrees, so the overall time is $O(n\log n)$.  We also build the Suffix array as a preprocess step, which takes an additional $O(n\log n)$ time.

As for space - the only non-trivially linear part of our solution is the array $A$ used in Batched Suffix Sort. Since we never let $se$ the number of sorted elements exceed $2n$, the lists in $A$ never contain more than $2n$ elements collectively. So the size of $A$ is always linear. After executing the Batched Suffix Sorting, we iterate the sorted lists to count the $P$-light occurrences and then reuse the space occupied by these lists as they are no longer required. \hfill\ensuremath{\Box}

We are left with the task of counting the $P$-heavy occurrences of every leaf $z$. Consider a heavy path $P=u_1,u_2 \ldots u_x$. Our key sub-task for finding all the $P$-heavy occurrences of all the leaves in the sidetrees of $P$ is constructing a sorted list of the $P$-heavy words of the leaves.

Note that unlike $P$-light occurrences, $P$-heavy occurrences of a leaf $z$ of a sidetree emerging from $u_i$ can not be in a sidetree emerging from $u_i$. However, they must be leaves of a sidetree emerging from $u_j$ for some $j > i$.

The process of building the sorted list of $P$-heavy words relies on the same principles we used for the $P$-light occurrences. However, there is a further difficulty to tackle. With $P$-light occurrences that lie on the same $u_i$ - we have a guarantee that the words match until the index $|w_i|$. Therefore, it is sufficient to sort by the suffixes starting right after the mismatch in index $|w_i| +2$. With the $P$-heavy words, we may have to compare $P$- heavy words from sidetrees of different nodes $u_i$ and $u_j$ with $i < j$. In this case, there is no guarantee that the words match in the indices in $[|w_i + 1| \ldots |w_j|]$.

To handle this difficulty, we partition the leaves into classes prior to sorting them. Our partition will have the property that the $P$-heavy words of leaves in the same class have a certain common prefix that exceeds the index in which the error occurs ($c$ is replaced by $d_i$). This property will allow us to sort the $P$-heavy words in every class using Batched Suffix Sorting.

The first step for sorting the $P$-heavy words is to partition the leaves in the sidetrees of $P=u_1,u_2 \ldots u_x$ by the $lcp$ of their $P$-heavy words with $w_x$. This is done with the following procedure:

{\bf LCP Partition:}\\
\fbox{\begin{minipage}{12cm}

{\sf Input:} A heavy path $P = u_1,u_2 \ldots u_x$
\\\\
{\sf Initialize} an array $LCP[1\ldots m]$ of size $m$ of empty lists. Let $j_x$ be an index in which $w_x$ occurs.
\\\\
\textbf{Alignment Step:}
{\sf For every $i \in [1\ldots x]$:}
For every leaf $z$ in a sidetree emerging from $u_i$:
\begin{enumerate}
    \item Extract an index $j_z$ in which $\mathcal{L}(z)$ occurs in $S$.
    \item \label{step:findlz} Find $l_z = lcp(hw(z),w_x)$ by computing $l_z = min(|w_i| + 1 + lcp(j_z + |w_i| + 2,j_x + |w_i| + 2),|w_x|)$.
    \item Compare between the symbols in index $l_z+1$ in $hw(z)$ and in $w_x$ in order determine the lexicographical order $o_z \in \{ < ,> ,=\}$ between $hw(z)$ and $w_x$ (For example, $o_z = <$ if $hw(z) <_L w_x$). If $l_z = |w_x|$, $o_z$ is set to $'='$.
    \item Add the tuple $(z,o_z)$ to $LCP[l_z]$.
    \end{enumerate}

\textbf{Insertion Step:}
\\\\
{\sf For every $l \in [1\ldots |w_x| -  1]$ (in increasing order):}
\begin{enumerate}
    \item If the list $L = LCP[l]$  is empty - do nothing.
    \item Otherwise, create 2 lists $L_l^>$ and $L_l^<$.
    \item For every tuple $(z,o_z)$, add $z$ to $L_l^{o_z}$.
\end{enumerate}

{\sf If} $L=LCP[|w_x|]$ is not empty, construct a new list $L_{|w_x|}$ and add $z$ to $L_{|w_x|}$ for every pair $(z,=) \in L$.
\end{minipage}}

Note that $l_z \ge |w_i| + 1$ since $hw(z)[1 \ldots |w_i|] = w_i= w_x[1\ldots |w_i|]$ and $hw(z)[|w_i + 1|] = d_i$. With that observation, it is clear that the formula for finding $l_z$ in Step \ref{step:findlz} works.

We make the following observation:

\begin{observation}\label{obs:lcplists}
For every list $L_l^>$ (or $L_l^<$ or $L_{|w_x|}$), every vertex $z\in L_l^>$ has $lcp(hw(z),w_x) = l$ and $hw(z)[l+1 \ldots m] = \mathcal{L}(z)[l+1 \ldots m]$. The running time of LCP partition is $O(m + |SE|)$ with $SE$ being the set of leaves in sidetrees of $P$.
\end{observation}
\textbf{Proof:} The $lcp$ property is derived directly from the construction of $LCP[1 \ldots m]$. As for complexity, every leaf is processed with a single $lcp$ query and a constant amount of basic operations. The iteration and construction of $LCP$ takes an additional $O(m)$ \hfill\ensuremath{\Box}

It follows from Observation \ref{obs:lcplists} that the lexicographical order between the $P$-heavy words of the vertices in $L_l^>$ (or $L_l^<$ or $L_{w_x}$) are determined by the suffixes starting in index $l+1$ of $L(z)$. Therefore, sorting $L_l^>$ by the lexicographical order of the heavy words can be done using the algorithm Batched Suffix Sorting.

\begin{theorem}\label{t:fastphwt}
The lexicographically sorted list of $P$-heavy words of a heavy path $P$ can be evaluated in $O(|SL|)$ amortized time with $SL$ the set of leaves in the sidetrees of $P$
\end{theorem}
\textbf{Proof:} We want to use LCP Partition on a batch of heavy paths. Transforming LCP Partition to a batched algorithm can be done with the same technique that was used to generate Batched Suffix Sorting from Suffix Sorting.

As in the Batched Suffix Sorting algorithm, we execute the alignment step of LCP Partition for possibly multiple heavy paths $P$ until the collective amount of leaves considered is between $n$ and $2n$. Once this amount is met, we construct the lists $L^>_l$, $L^<_l$ and $L_{|w_x|}$ for all the paths in the batch by applying the insertion step. We then sort the lists by the lexicographic order of $hw(z)$ with Batched Suffix Sorting. The overall time is $O(n + m) = O(n)$.

We are left with the task of merging the sorted lists $L_l^>$,$L_l^<$ and $L_{|w_x|}$ into a single sorted list $L$ containing all the $P$-heavy words. This is done by applying the following observation:

\begin{observation} \label{obs:concsort}
Let $l_1,l_2, \ldots l_c$ be the set of indices for which either $L_{l_i}^<$ or $L_{l_i}^>$ is constructed for the path $P$ by LCP Partition. The sorted list $L$ of the $P$-heavy words is of the form $L= L_{l_1}^< , L_{l_2}^<, \ldots L_{l_c}^<,L_{|w_x|},L_{l_c}^>,L_{l_{c-1}}^> \ldots L_{l_1}^>$. (If $L_{l_i}$ was not constructed, it is considered as an empty list)
\end{observation}
\textbf{Proof:} We start by showing that the lists $L_{l_i}^<$ must appear in increasing order of $l_i$ in $L$. Let $a,b \in \{ l_1,l_2 \ldots l_c\}$ be two indices for which $a<b$. Let $w_a \in L_a^<$ and $w_b \in w_b^<$. Since the LCP of $w_b$ and $w_x$ is $b \ge a + 1$, we have $w_b[a+1] = w_x[a+1]$. Since the LCP of $w_a$ and $w_x$ is $a$ and $w_a <_L w_x$, we have $lcp(w_a,w_b) = a$ and $w_a[a+1] <_L w_x[a+1] = w_b[a+1]$ and therefore $w_a <_L w_b$. Similar arguments can be made to prove that $w_b <_L w_a$ for every $w_a \in L_{l_i}^>$ and $w_b \in L_b^>$. It is straight forward from the construction of the lists $L_{l_i}^<$ and $L_{l_i}^>$ that for every $l_i$ and every $W \in L_{l_i}^>$ , $w \in L_{l_i}^<$ and $w' \in L_{|w_x|}$ we have $w <_L w' <_L W$. \hfill\ensuremath{\Box}

With Observation \ref{obs:concsort}, the construction of the sorted $P$-heavy words list is completed. Observe that LCP Partition naturally generates $L_l^>$ and $L_l^<$ in increasing order of $l$. Therefore, the concatenation of the lists in the order dictated by Observation \ref{obs:concsort} does not require any further sorting and can be executed in linear time, and the proof of Theorem \ref{t:fastphwt} is completed.\hfill\ensuremath{\Box}

Given the sorted list $L_P[1 \ldots h]$ of the $P$-heavy words, we are interested in finding the node $z'$ with $L(z') = hw(z)$ for every word $hw(z) \in L_P$. We can do this in linear time as follows: First, observe that every $P$-heavy word has the prefix $w_1$. So $z'$ ,if it exists, must be a descendant of $u_1$ and therefore is a leaf in a sidetree of $P$. Let $L[1\ldots l]$ be the sorted list of occurrences of $w_1$ stored in $u_1$. These are actually all the leaves in the sidetrees of $P$. The following procedure matches every $hw(z) \in L_P$ with its corresponding $z'$:
\\\\
{\bf Count $P$-Heavy:}\\
\fbox{\begin{minipage}{12cm}

{\sf Input:} The lexicographically sorted lists $L_P[1 \ldots h]$ of $P$-heavy words and $L[1 \ldots l]$ the list of lexicographically sorted vertexes in the sidetrees of $P$
\\\\
{\sf Initialize} two indices $i = j = 1$.
\\\\
\textbf{While $i \le h$ and $j \le l$:}
\begin{enumerate}
    \item Let $hw(z) = L_P[i]$ and $\mathcal{L}(z') = L[j]$
    \item \textbf{If $hw(z)=\mathcal{L}(z')$:}
    \begin{enumerate}
        \item Increase the counter associated with $z$ by $Oc(z')$.
        \item Increase the counter associated with $z'$ by $Oc(z)$.
        \item Increase $i$ by $1$.
\end{enumerate}
    \item \textbf{If $hw(z) <_L \mathcal{L}(z')$:} Increase $i$ by $1$.
    \item \textbf{If $hw(z) >_L \mathcal{L}(z')$:} Increase $j$ by $1$.
\end{enumerate}
\end{minipage}}

It can be easily verified that Count $P$-Heavy counts the $P$-heavy occurrence $z'$ of every leaf $z$ in a sidetree of $P$. Notice that double counting will not occur. That is due to the following:

\begin{fact} \label{f:incij}
Let $z$ and $z'$ be two leaves in sidetrees of $P$ emerging from $u_i$ and $u_j$ respectively such that $hw(z) = \mathcal{L}(z')$. It must be the case that $i < j$.
\end{fact}
Fact \ref{f:incij} guarantees that if we count the occurrences of $z$ as $1$-ham occurrences of $z'$ and vice versa when $hw(z)$ and $\mathcal{L}(z')$ are be visited in Count $P$-Heavy, we will not count them as $1$-ham occurrences of each other again, because it can't be the case that $hw(z') = \mathcal{L}(z)$.
\\\\
The lexicographic comparisons between $hw(z)$ and $\mathcal{L}(z')$ can be executed in constant time using $lcp$ queries. To efficiently execute an $lcp$ query with a $P$-heavy word, we store the $P$-heavy word $hw(z)$ as a pair $(z,i)$ with $i$ the index in which $\mathcal{L}(z)$ is modified. With that representation, two $lcp$ queries can be used to find $a = lcp(hw(z),\mathcal{L}(z'))$ in a 'kangooroo' jump manner. If $a < m$, the following symbol can be compared to determine the lexicographic order between $hw(z)$ and $\mathcal{L}(z')$. With the constant time lexicographic comparing, it is easy to see that the complexity of Count $P$-Heavy is $O(|L_P| + |L|) = O(|SE|)$ with $SE$ being the set of leaves in the sidetrees of $P$.

Note that when the equality $hw(z) = L(z')$ is met, it is crucial to increase $i$ rather than $j$. That is due to the fact that $L_P$ may contain duplicates while $L$ does not. Alternatively, $L_P$ can be preprocessed to group duplicates together. We conclude the counting of $P$-heavy occurrences with the following:
\begin{theorem}\label{t:fastheavy}
The $P$-heavy occurrences of every leaf $z$ can be counted over all the heavy paths $P$ such that $z$ is a leaf in a sidetree of $P$ in $O(n\log n)$ time and linear space.
\end{theorem}
\textbf{Proof:} For every heavy path $P = u_1,u_2 \ldots u_x$, we use Theorem \ref{t:fastphwt} to obtain the list $L_P$ of sorted $P$-heavy words and obtain $L$ from $u_1$. We then apply Count $P$-Heavy on $L_P$ and $P$ to match every $hw(z) \in L_P$ with its $P$-heavy occurrence $z'$ if exists, and update the corresponding counters accordingly.

The amortized time for applying Theorem \ref{t:fastphwt} for a path $P$ is $O(|SE(P)|)$ with $SE(P)$ being the set of leaves in the sidetrees of $P$. Every leaf in $ST$ is a leaf in the sidetree of at most $\log(n)$ heavy paths, so the overall complexity is $O(n\log n)$. We also construct the suffix array as a preprocess procedure, which takes an additional $O(n\log n)$ time.

As for space, the only non-trivially linear part is the array $LCP[1 \ldots m]$ used in LCP Partition. As before, we apply LCP Partition on batches of size at most $2n$, so the collective size of the lists in $LCP[1\ldots m]$ never exceeds $O(n)$. After obtaining $L_P$ for all the paths in the batch, we apply Count $P$-Heavy for every path in the batch and then reuse the space occupied by the sorted lists $L_P$ as they are no longer required. \hfill\ensuremath{\Box}

When put together, Theorem \ref{t:fastlight} and Theorem \ref{t:fastheavy} yield the main result of this section:
\begin{theorem}\label{t:1map}
The $1$-mappability problem can be solved using $O(n\log n)$-time and linear space on a text with infinite integer alphabet.
\end{theorem}

Note that for infinite integer alphabet, better time can not be achieved unless certain values of $m$ are excluded. For example:
\begin{observation}
For a text $S$ over infinite integer alphabet and $m=2$, there is an index $i \in [1\ldots n]$ with at least $1$-ham occurrence iff the symbols of $S \cdot \sigma'$ are not distinct for some $\sigma' \notin \Sigma$.
\end{observation}
The above straight forward observation shows a trivial relation between the $k$-mappability problem and reporting whether or not all the elements of a set are distinct - which can not be done in $o(n\log n)$. It can be easily generalized for every fixed value of $m$.

\section{$O({n^2\over m^2} + n)$ and Linear Space Algorithm for $1$-Mappability with Constant sized Alphabet}\label{s:lspace}
As a warm up, we present a technique for counting the $1$-ham occurrence of a word with size $m$ in $O(\frac{n}{m})$ time. Applying this technique to every $m$-sized word yields an $O(\frac{n^2}{m})$ algorithm for $1$-mappability. We then proceed to show how to process all the words of size $m$ not one by one, but in batches of size $O(m)$. We extend the technique used in the warm up to handle a batch in $O(\frac{n}{m} + m)$ time. Since there are $O(\frac{n}{m})$ batches, this yields an $O((\frac{n}{m})^2 + n)$ time algorithm.
\subsection{Warm up - $O(\frac{n^2}{m} + n)$}
Let $w= S[i \ldots i+m-1]$ be a subword of $S$ with length $m$.
 \begin{definition} \label{d:loccrocc}
 Let $w_1=S[j \ldots j+m-1]$ be a $1$-ham occurrence of $w$. $w_1$ is an $l$-occurrence of $w$ if $w_1[1\ldots \lceil \frac{m}{2} \rceil ] = w[1\ldots \lceil \frac{m}{2} \rceil]$. $w_1$ is an $r$ occurrence of $w$ if $w_1[\lceil \frac{m}{2} \rceil +1 \ldots m] = w[\lceil \frac{m}{2} \rceil +1 \ldots m]$. We respectively denote as $Lo(w)$ and $Ro(w)$ the sets of $l$-occurrences and $r$-occurrences of $w$ in $S$.
 \end{definition}

It is easy to see that $|Lo(w)| + |Ro(w)| - \#w$ is the number of $1$-ham occurrences of $w$, with $\#w$ denoting the number of proper occurrences of $w$ in $S$. In this section, we show how to evaluate the number of $l$-occurrences of a given word $w$ in $O(\frac{n}{k})$ time. A symmetrical approach can be applied to count the number of $r$-occurrences of $w$. $\#w$ can be evaluated for all the subwords of $S$ in $O(n)$ time using the suffix tree.

\begin{theorem}\label{t:apRep}
All the occurrences of a string $w$ of size $m$ in a text of size $n$ can be represented by a set of $O(\frac{n}{m})$ arithmetic progressions of the form $A = (s,e,d)$ such that $A = (s,e,d)$ represent a sequence of occurrences with starting indexes $\{i_x = s + d\cdot x| x\ge 0 ,i_x \le e \}$. If $w$ is periodic, every arithmetic progression $A=(s,e,d)$ has $d =per(w)$. $|A|=e-s+1$ represents the number of occurrences represented by $A$. Every arithmetic progression that has $A > 1$  corresponds to a periodic set of instances contained within a run with period $d$. This representation is called the periodic occurrences representation of $w$ and it can be obtained in $O(\frac{n}{m})$ time from the suffix tree following an $O(n)$ time preprocessing.
\end{theorem}

A proof for the above can be found in Section \ref{A:complspace}.

Given a words $w=S[i \ldots i +  m- 1]$, we use Theorem \ref{t:apRep} to obtain all the occurrences of $w_L = w[1 \ldots \lceil \frac{m}{2} \rceil]$ in periodic occurrences representation. For every occurrence of $w_L$ in this representation, we wish to check if it is a prefix of an $l$-occurrence.

We process every arithmetic progression $A = (s,e,d)$ of occurrences of $w_L$. If $A$ only represents a single occurrence of $w_L$ in index $s$, we query $l_1 = LCP(s,i)$. If $l_1 \ge m$, we have a proper instance of $w$. Otherwise, we have a mismatch. Proceed to query $l_2 = LCP(s + l_1 + 1, i + l_1 + 1)$. If it is the case that $l_2 + l_1 + 1 \ge m$, we count $s$ as an $l$-occurrence of $w$.

If $A=(s,e,d)$ represents multiple occurrences of $w_L$, then $w_L$ must have a period $d$. We exploit the periodic structure of the occurrences represented by $A$ to compute $l_2$ for all the occurrences in $A$ using constant time. The following lemma proven in Section \ref{A:complspace} is the key for doing so.

\begin{lemma} \label{o:RepetativeLCP}
Let $A=(s,e,d)$ be an arithmetic progression representing a set of indexes $s_j = s + j \cdot d$ for $j \in [0 \ldots |A|-1]$ within a run with period $d$.

Let $i \in [1\ldots n]$ be an index and let $l_p = lcp(i,s)$. Let $Ex_i$ be the maximal extension of a run with period $d$ containing $i$ to the right of $i$ (regardless of periodicity, $Ex_i\ge d$), and let $Ex_s$ be the maximal extension of the period $d$ to the right of $s$.

\begin{enumerate}
    \item If $l_p < d$:  $LCP(i,s_j) = l_p$ for every $j \in [0 \ldots |A|-2]$.
    \item Otherwise, $LCP(i,s_j) = min(Ex_i,Ex_s - j\cdot d)$  for every $j\in [0 \ldots |A|-1]$ such that $Ex_i \neq Ex_s - j \cdot d$.
\end{enumerate}

\end{lemma}

We exploit Lemma \ref{o:RepetativeLCP} to efficiently implement the following subroutine (details proof for the following can be found in Section \ref{A:complspace}.
\begin{lemma}\label{l:fastreplcp}
Given an arithmetic progression $A=(s,e,d)$ representing the indexes $\{s_j = i +j\cdot d|j \in [0\ldots |A|-1]\}$ that are contained within the same run with period $d$, and an index $s\in [1\ldots n]$. The values $lcp_j=LCP(s,s_j)$ can be evaluated and represented in $O(1)$ following $O(n)$ preprocess time on $S$.
\\\\
The representation consists of pairs $(I,L)$ such that $I=[a\ldots b]$ is a consecutive interval of $j$ values and $L$ is an integer such that one of the following holds:
\begin{enumerate}
    \item \label{type:const} $lcp_j=L$ for every $j\in I$
    \item \label{type:arit} $lcp_j = L -j\cdot d$ for every $j\in I$.
\end{enumerate}
Every pair is stored alongside with a bit indicating which one of the above holds for this pair.
\end{lemma}

In the process of evaluating the representation of $lcp_j$ for $A=(s,e,d)$ and $i$, at most one of the indexes in $A$ is called the aligned index. In the case in which $l_p < d$, the aligned index is $|A|-1$. In the case in which $l_p \ge d$, $j^*$ such that $Ex_s - j^* \cdot d$ is the aligned index, provided that it is an integer. We mark the pair representing the LCP value of the aligned index.

We employ Lemma \ref{l:fastreplcp} to obtain a representation of $l_1^j$ for every $j\in [0 \ldots |A|-1]$. After obtaining this representation, we are left with the task of applying a second $LCP$ query after the mismatch index for every $s_j$ (That will be the equivalent of finding $l_2$). Namely, for every $s_j$ we need to compute $l_2^j= LCP(i+l_1^j + 1,s_j + l_1^j + 1)$. More precisely , we need to count the number of $j$ values for which $l^j= l_1^j + l_2^j + 1 \ge m$.

For every pair $(I = [a \ldots b] , L)$, we wish to evaluate $l^j_2$ for $j \in I$ by employing Lemma \ref{l:fastreplcp} again. In order to do that, we first need to prove that the settings of Lemma \ref{l:fastreplcp} are satisfied in the second evaluation. We prove the following lemma in Section \ref{A:complspace}.

\begin{lemma}\label{l:conditionforfast}
For every pair $(I,L)$ in the output of Lemma \ref{l:fastreplcp} on $A=(s,e,d)$ and $i$ that is not corresponding to an aligned occurrence, one of the below holds for $j\in I$.
\begin{enumerate}
    \item \label{condition:1} $i+lcp_j+1$ is a fixed value and $s_j + lcp_j+1$ is an arithmetic progression of indexes within a run with period $d$ with difference $d$.
    \item \label{condition:2} $i+lcp_j+1$ is an arithmetic progression of indexes within a run with period $d$ with difference $d$ and $s_j + lcp_j+1$ is a fixed value
\end{enumerate}
\end{lemma}

It follows from Lemma \ref{l:conditionforfast} that Lemma \ref{l:fastreplcp} can be applied to each of the pairs representing the non aligned indexes to evaluate a representation of $lcp^j_2$ for every non aligned index $j$ in $O(1)$ time. The aligned index, if exists, has its $l_2^{j}$ evaluated individually.

The above process outputs a set of (at most) 4 non-singular intervals for which the values of $l_2^j$ and $l_1^j$ are represented either as an arithmetic progression or as a fixed value. We can easily deduce the amount of occurrences $s_j$ with $l_1^j + l_2^j + 1 \ge m$ from this representation. We sum the amount of occurrences of $w_L$ with $l_1^j + l_2^j + 1 \ge m$ over all the arithmetic progressions of occurrences to obtain $|Lo(w)|$. A symmetric procedure can be constructed to evaluate $|Ro(w)|$ using occurrences of $w_R = w[\lceil \frac{m}{2} + 1 \rceil \ldots m]$. We can use the suffix tree to obtain $\#w$ (the number of occurrences of $w$ within $S$) for every $m$-length word in $S$ in $O(n)$ time.
The number of occurrences of $w$ with at most one mismatch is $|Ro(w)| + |Lo(w)| - \#w$. substracting $\#w$ is required to omit double counting. We do this process for every word of size $m$.

\textbf{Complexity.} For a word $w$ of size $m$, we process the arithmetic progressions of occurrences of $w_L$. For every arithmetic progression $A$, we evaluate a representation of $lcp^j_1$ and $lcp^j_2$ in constant time using Lemma \ref{l:fastreplcp}. We deduce the number of $l$-occurrences corresponding to the occurrence of $w_L$ represented by $A$ from the representation of $lcp^j_1$ and $lcp^j_2$ in constant time. We execute a symmetric procedure to deduce the number of $r$-occurrences of $w$ as well. There are $O(\frac{n}{m})$ arithmetic progressions in periodic occurrences representation of $w_L$, so counting the 1-ham occurrences of a single word takes $O(\frac{n}{m})$. We do this for every word of size $m$, so it adds up to $O(\frac{n^2}{m})$. There is an additional $O(n)$ preprocessing time prior to the iteration on the words to enable $LCP$ queries, suffix tree construction and access to the periodic occurrences representation. The overall complexity is $O(\frac{n^2}{m} + n) = O(\frac{n^2}{m})$.

\subsection{Reducing the complexity to $O(\frac{n^2}{m^2} + n)$}
For reducing the complexity by a factor of $m$, we present a technique for obtaining  $|Lo(w)|$ for a batch containing $O(m)$ words in $O(\frac{n}{m}+m)$ time. Consider the consecutive set of words with length $m$ starting in the indices $[i \ldots i + \frac{m}{4}]$. For every word $w^t = S[i + \frac{m}{4} - t \ldots i + \frac{m}{4} - t + m - 1]$ with $t \in [0 \ldots \frac{m}{4}]$ in this set, the left half of $w^t$ denoted as $w^t_L = S[i + \frac{m}{4} - t \ldots i + \frac{m}{4} - t + \lceil \frac{m}{2} \rceil]$ contains the word $w^i_L = S[i + \frac{m}{4} \ldots  i +\frac{m}{2} - 1]$. We use the occurrences of $w^i_L$ to evaluate $|Lo(w^t)|$ for every $t \in [0 \ldots \frac{m}{4}]$ as we did in the previous section. A symmetric process can be constructed for computing $|Ro(w)|$.

We start by finding the arithmetic progression representation of the occurrences of $w^i_L$. For simpler notation, we denote the starting and ending indices of $w^i_L$ as $w^i_L = S[s_i \ldots e_i]$. For every cluster $A$ with occurrences $\{s_j = s + d \cdot j| j \in [0 \ldots |A|-1] \}$, let $r_1^j = lcp( s_j, s_i)$ and $r_2^j =lcp (s_j + r_1^j + 1, s_i + r_1^j + 1)$. As in the previous section, we use Lemma \ref{l:fastreplcp} to obtain a compact representation of $r^j = r_1^j + r_2^j + 1$ for every index $s_j$ represented by $A$.

Using the following Lemma ,that can be proved similarly to Lemma \ref{l:fastreplcp}, we obtain a compact representation of $l^j = LCS(s_i - 1,s_j - 1) $ for every $j\in [0 \ldots |A|-1]$.
\begin{lemma}\label{l:fastreplcs}
Given an arithmetic progression $A=(s,e,d)$ representing the indexes $\{s_j = i +j\cdot d|j \in [0\ldots |A|-1]\}$ that are contained within the same run with period $d$, and an index $s\in [1\ldots n]$. The values $lcs_j=LCS(s,s_j)$ can be evaluated and represented in $O(1)$ following $O(n)$ preprocess time on $S$.
\\\\
The representation consists of pairs $(I,L)$ such that $I=[a\ldots b]$ is a consecutive interval of $j$ values and $L$ is an integer such that one of the following holds:
\begin{enumerate}
    \item \label{type2:const} $lcs_j=L$ for every $j\in I$
    \item \label{type2:arit} $lcs_j = L + j\cdot d$ for $j\in I$.
\end{enumerate}
\end{lemma}

After obtaining the values of $l_j$ and $r_j$, our next task is deducing for every $s_j$ represented by $A$, what are the values of $t$ for which $s_j$ is corresponding to an $l$-occurrence of $w^t$. The following observation is the key for doing so.
\begin{observation}\label{o:perContribution}
$s_j - t$ is an $l$-occurrence of $w^t$ iff $r^j \ge m - t$ and $l^j \ge t$.
\end{observation}

Observation \ref{o:perContribution} allows us to associate every occurrence $s_j$ with a continuous interval $I=[a\ldots b]$ such that $s_j$ is extendable to an $l$-occurrence of $w^t$ for and only for $t\in I$.

We initialize a data structure $D$ for maintaining $\frac{m}{4}$ counters $C_0,C_1,C_2 \ldots C_{\frac{m}{4}}$. $C_t$ counts $l$-occurrences of $w^t$. Initially, $C_t = 0$ for every $t \in [0\ldots \frac{k}{4}]$. We already know from observation \ref{o:perContribution} that for every $s_j$, the indexes $C_t$ with $t\in [m-r_j \ldots l_j]$ need to be increased by $1$. We call this type of updates, in which a consecutive interval of counters is increased by a constant value, an interval increment update. There are folklore techniques for applying this kind of updates to an array of counters efficiently.

Unfortunately, an efficient data structure for applying interval increment updates will not be sufficient for our cause, as we wish to process the effect of a \textbf{set} of occurrences on $D$. We therefore need to explore the structure of the set of updates $[m -r_j \ldots l_j]$ derived from the occurrences $s_j$ represented by a cluster $A$.

We present the following types of updates to be applied to an array of counters $D$.
\begin{definition}
An interval increment is represented by a triplet $(a,b,x)$. Applying $(a,b,x)$ to $D$ results in every counter $C_t$ with $t\in [a\ldots b]$ being increased by $x$.

An {\em increasing stairs update} is represented by a triplet $(a,b,p)$. The update requires applying the following modifications on $D$:

For every $d\in 1 \ldots \lfloor \frac{b-a+1}{x} \rfloor$ Counters $C_i$ with $i \in [a + p \cdot (d-1) \ldots a+ p \cdot d - 1]$ are increased by $d$. The counters with $t\in [a + p \cdot \lfloor \frac{b-a+1}{x} \rfloor \ldots b]$ are increased by $\lfloor \frac{b-a+1}{x} \rfloor +1$

A {\em decreasing} stairs update is also represented by a triplet $(a,b,p)$. The update requires applying the following modifications on $D$:

For every $d\in 1 \ldots \lfloor \frac{b-a+1}{x} \rfloor$ Counters $C_i$ with $i \in [b - p \cdot d +1 \ldots b - p \cdot (d - 1)]$ are increased by $d$. The counters with $t\in [a \ldots b - p \cdot \lfloor \frac{b-a+1}{x} \rfloor]$ are increased by $\lfloor \frac{b-a+1}{x} \rfloor +1$

We call the interval $[a \ldots b]$ the {\em span} of the stairs. We call the interval that is increased by $d$ the $d$th step of the update. $p$ is called the {\em width} of the stairs update.

A {\em negative} stairs update (either increasing or decreasing) is a stairs update in which the counter in the $d$th step is decreased by $d$ rather than being increased by $d$.
\end{definition}

\begin{example}
Let $x = 10$ and an array $D = (0,0,0,0,0,0,0,0,0,0)$. An increasing stairs update $(3,8,2)$ on $D$ will result in the counters being set to $(0,0,1,1,2,2,3,3,4,0,0)$. Applying a decreasing stairs update $(1,5,2)$ on the updated counters in $D$ will result in the counters being set to $(3,2,3,2,3,2,3,3,4,0,0)$.
\end{example}

It turns out that a constant number of interval updates and stairs updates can be used to express the updates derived from the occurrences $s_j$ represented by a cluster $A$. In Section \ref{A:complspace}, we prove the following:
\begin{lemma}\label{l:representwithstairs}
Given a cluster $A$ of occurrences of $w^i_L$, the set of updates that need to be applied to $D$ in order to represent the $l$-occurrences corresponding to occurrences $s_j$ with $j\in [0\ldots |A|-1]$ can be represented by a constant number of stairs updates and interval increment updates. Given $A$ and the representation of $l_j$ and $r_j$, this set of stairs and interval increment updates can be retrieved in $O(1)$ time.

Over all the clusters representing occurrences of $w^i_L$, every stairs update $(a,b,p)$ in the representation has the same stairs width $p$ which is the period of $w^i_L$.
\end{lemma}

Our algorithm runs as follows: Initialize a data structure $D$ for maintaining a set of $\frac{m}{4}$ counters. Find all the occurrences of $w^i_L$ in arithmetic progression representation. For every one of the $O(\frac{n}{m})$ arithmetic progressions, find the arithmetic progressions representing $r^j$ and $l^j$. Apply Lemma \ref{l:representwithstairs} to obtain an $O(1)$ size set of interval increment update and stairs update that represents the required modifications to be applied to $D$. The final ingredient for our algorithm is a data structure that enables the efficient application of these updates. In Section \ref{s:stairs}, we prove the following.
\begin{theorem}\label{t:stairs}
An array of $t$ counters can be maintained to support stairs updates in $O(1)$ time per update. Retrieving the values of all the counters in the array takes $O(t+u)$ time with $u$ being the amount of applied updates. The data structure works in the restricted settings in which every update $(a,b,p)$ has the same $p$ value.
\end{theorem}

Note that the restriction on the queries hold in our case, since the step width is always $p$ the period of $w^i_L$ in all the stairs updates constructed in Lemma \ref{l:representwithstairs}.

Every update corresponding to a set of occurrences of a certain type is applied to $D$ in $O(1)$ by employing the data structure of Theorem \ref{t:stairs} all the updates take $O(\frac{n}{m})$ by applying Theorem \ref{t:stairs}.

Note that we need a data structure for handling interval increment updates with the same complexities as the data structure of Theorem \ref{t:stairs}. The construction of such a data structure is quite simple and may be considered folklore. We therefore omit the implementation details of this data structure.

After applying the updates, we query our data structure for the values of all the counters. This process takes $O(m + \frac{n}{m})$ time. This is done for batches of $\frac{m}{4}$ consecutive indices. The indices of $S$ are partitioned to $4\frac{n}{m}$ such batches. We also preprocess the text for constant time $lcp$ and $lcs$ queries and construct the suffix tree. The total running time is $O(n + \frac{n}{m}(m + \frac{n}{m})) = O(n + \frac{n^2}{m^2})$.  Recall that we described a procedure for evaluating $|Lo(w)|$. A symmetric procedure can be constructed to evaluate $|Ro(w)|$.

Note that for $m\in \Omega(\sqrt{n})$, $O(\frac{n^2}{m^2} + n)$ is dominated by $O(n)$. The main result of this section immediately follows.
\begin{theorem}
For constant size alphabet and $m \in \Omega(\sqrt{n})$, the $1$-mappability problem can be solved in time $O(n)$.
\end{theorem}

\section*{Acknowledgements}
We warmly thank Tomasz Kociumaka for useful discussions.

\bibliographystyle{plainurl}


\section{Introduction}

Task-oriented dialog systems aim to satisfy user goals pertaining to certain tasks, such as booking flights \citep{hemphill-etal-1990-atis}, providing transit information \citep{raux2005let}, or acting as a tour guide \citep{budzianowski2018multiwoz}. Neural models for task-oriented dialog have become the dominant paradigm \citep{williams2016end,wen2016network,zhao2017generative}. These data-driven approaches can potentially learn complex patterns from large dialog corpora without hand-crafted rules. However, the resulting models struggle to generalize beyond the training data and under-perform on unseen dialog tasks and domains \citep{zhao2018zero,rastogi2020towards}.

A long-standing challenge in dialog  research is to flexibly adapt systems to new dialog domains and tasks \citep{zhao2018zero,mosig2020star}. Consider a system that has been trained to handle several different tasks (e.g., restaurant reservations, ride booking, weather, etc.). How can this dialog system be \textit{extended to handle a new task} (e.g., hotel booking), without collecting additional data? This paper tackles this challenge and aims to address the problem of zero-shot generalization using the \textbf{schema-guided paradigm}.

\begin{figure}
    \centering
    \includegraphics[width=0.5\textwidth]{schema_paradigm.png}
    \caption{In the standard paradigm, data driven models implicitly learn the task-specific dialog policies (i.e., schemas). This precludes generalization to an unseen task at inference time. In contrast, in the schema-guided paradigm, dialog policy is explicitly provided to the model through a schema graph. At inference time, the model is given the schema for the new task and can therefore generalize in a zero-shot setting. }
    \label{fig:schema_paradigm}
\end{figure}
The advent of large-scale pre-training \citep{devlin-etal-2019-bert,radford2019language} has led to significant progress in domain adaptation across areas in NLP, including natural language understanding \citep{wang2018glue,wang2019superglue}, open-domain dialog \citep{zhang-etal-2020-dialogpt,adiwardana2020towards} and language understanding for task-oriented dialog \citep{wu2019transferable,mehri2020dialoglue}. Generalization in end-to-end task-oriented dialog has proven to be significantly more difficult, particularly in zero-shot settings where there is no training data \citep{mosig2020star}. We posit that it is inherently difficult to generalize to unseen dialog tasks because of the \textbf{dialog policy}.

Traditionally, an end-to-end dialog system must perform three distinct tasks. First, it must understand the dialog history and identify any relevant user intents or slots. Next, it must decide on the appropriate system action, according to a task-specific dialog policy. Finally, it must generate a natural language utterance corresponding to the system action. In a pipeline dialog system, these three steps are performed by the NLU, DM and NLG respectively \citep{jurafsky2000speech}. Prior work has exhibited generalizability in language understanding and, to a lesser extent, in language generation. However for end-to-end dialog, the task-specific dialog policy inherently precludes zero-shot generalization. An end-to-end dialog model trained on several tasks, will \textit{implicitly} learn the dialog policies from the data. However, when generalizing to a new task in a zero-shot setting, the model has no knowledge of the dialog policy for the \textit{new} task.

To address the difficulty of generalizing to new task-specific dialog policies and in order to facilitate zero-shot generalization, we present the \textbf{schema-guided paradigm}. Generally, end-to-end neural models implicitly learn the task-specific dialog policies from large corpora. In contrast, in the schema-guided paradigm, we explicitly provide the task-specific dialog policies to the model in the form of a \textbf{schema graph}. The schema graphs define the system's behavior for a specific task (e.g., when the user provides the reservation time, ask them for the number of people). When transferring to an unseen task, the corresponding schema graph is explicitly provided to the model. This enables language understanding and the dialog policy to be decoupled. The model no longer needs to implicitly memorize the task-specific policies from the training data. Instead, the model learns to interpret the dialog history and align it to the schema graph. As such, when transferring to a new task, the schema graph serves as an inductive bias that provides the model with the task-specific dialog policy.

To address the challenge of zero-shot transfer learning, \citet{mosig2020star} presented the STAR corpus and several baseline experiments. We extend their baselines for the task of \textit{next action prediction}. We introduce the \textbf{Schema Attention Model} (\textsc{sam}) and thorough schema representations for the 24 different tasks in the STAR dataset. \textsc{sam} obtains a \textbf{+22 $\mathbf{F_1}$} score improvement over baseline approaches in the zero-shot setting, validating the schema-guided paradigm and demonstrating the feasibility of zero-shot generalization for task-oriented dialog. Our code and model checkpoints are open-sourced and be found at \textbf{\url{https://github.com/shikib/schema_attention_model}}.

\section{Related Work}

\subsection{Zero-Shot Dialog}

Zero-shot transfer learning has been of interest to the dialog research community. Many approaches have been proposed for zero-shot adaptation of specific dialog components. \citet{chen2016zero} present a zero-shot approach for learning embeddings for unseen intents. \citet{bapna2017towards} show that slot names and descriptions can be leveraged to implicitly align slots across domains and achieve better cross-domain generalization. \citet{wu2019transferable} similarly use slot names, in combination with a generative model for state tracking, to obtain strong zero-shot results. \citet{shah2019robust} leverage examples for zero-shot slot filling. Generally, approaches for zero-shot generalizability leverage the similarity across domains (e.g., \textit{restaurant-area} and \textit{hotel-area} are conceptually similar). The advent of large-scale pre-training \citep{devlin-etal-2019-bert,radford2019language} allows for language understanding across dissimilar domains. \citet{rastogi2020schema} address zero-shot domain adaptation in state tracking by leveraging BERT \citep{devlin-etal-2019-bert} with a domain-specific API specification.

\citet{zhao2018zero} present an approach for zero-shot end-to-end dialog. They leverage the Action Matching framework to learn a cross-domain latent action space. \citet{qian2019domain}  use model-agnostic meta learning to attain stronger results in zero-shot dialog. Both these approaches rely on additional annotations, which make them unsuitable for the STAR corpus. While there is a significant amount of work in zero-shot generalizability for language understanding, there is considerably less research in adaptation for end-to-end dialog\footnote{While we focus on next action prediction, in the STAR dataset it is trivial to go from a system action to a natural language response and as such we consider our task to be end-to-end dialog.}. This is in part because of the difficulty of generalizing to unseen task-specific policies. To this end, \citet{mosig2020star} presented STAR, a corpus consisting of 24 different dialog tasks, and several baseline models for zero-shot adaptation on STAR. The results in this paper significantly outperform the baselines introduced by \citet{mosig2020star} as we leverage the schema-guided paradigm for zero-shot generalizability in dialog.

\subsection{Schema-Guided Paradigm}

Plan-based dialog systems \citep{ferguson1998trips,rich1998collagen,bohus2009ravenclaw} reason about user intent, in the context of a \textit{dialog plan}. RavenClaw \citep{bohus2009ravenclaw} consists of a task specification that defines the behavior of a system depending on various user actions. Plan-based dialog systems decouple the task-specific dialog policy from the task-agnostic components of the system. This allows a system to be extended to a new task by updating the task specification. The schema-guided paradigm shares a similar motivation, and aims to disentangle the dialog policy in neural, data-driven dialog systems.

Several approaches have been presented to discover dialog structure graphs (similar to the schemas in this paper) from data in an unsupervised manner \citep{shi2019unsupervised,qiu2020structured,xu2020discovering,hu2019gsn}. These approaches have been used to enhance generation for open-domain dialog \citep{qiu2020structured,hu2019gsn}. To the best of our knowledge, these dialog structures have neither been used for generation in task-oriented dialog nor in zero-shot settings. While our schemas are similar to these structure graphs, they are hand-crafted similar to those in plan-based dialog systems. Future work may extend our work by leveraging unsupervised structure graph discovery as an alternative to hand-crafted schemas.

\section{Task Definition}

We address the problem of transferring dialog models to unseen tasks and domains \citep{zhao2018zero}. This problem is especially important in real world settings. It is impossible to preconceive every dialog task that users may need (e.g., a COVID-19 information dialog system). Furthermore, collecting new dialog data for each new task is inherently unscalable \citep{rastogi2020towards}. While rule-based/pipeline dialog systems may be easier to extend to new tasks \citep{bohus2009ravenclaw}, there is a tradeoff between the adaptability of non-neural systems and the performance of neural models.

\subsection{STAR Dataset}

The STAR dataset \citep{mosig2020star} was collected for the purpose of studying transfer learning in dialog. The dataset spans 24 different tasks in 13 different domains (e.g., the restaurant domain has \textit{`restaurant-search'} and \textit{`restaurant-reservations'}). The data collection procedure was designed to reduce ambiguity in the system responses and make system actions deterministic. As such, Amazon Mechanical Turk (AMT) workers were given a flow chart diagram for each task. This flow chart defined the task, including the order in which questions should be asked (e.g., ask date before city), how to respond to various user responses and how to query a database. Additionally, in order to minimize variance in the responses from the wizard, \citet{mosig2020star} incorporate a \textit{suggestions module} during data collection. This module maps the wizard utterance to the closest pre-written response (e.g., \textit{`Give me your name'} $\rightarrow$ \textit{`What is your name?'}). In some cases, it is not possible for the AMT worker to use the suggestions module. Nonetheless, the module increases the consistency of the system actions.

\citet{mosig2020star} present baseline results on the tasks of next action prediction and response generation. The present paper focuses on \textit{next action prediction}. The objective of next action prediction is to predict the correct system action conditioned on the dialog history. Since there is a one-to-one mapping between system actions and corresponding natural language responses, the primary challenge in extending a next action prediction model to response generation resides in learning to accurately fill in the response templates (e.g., \textit{`Your reservation is confirmed for \{date\}'}).

The STAR dataset consists of three different types of dialogs: (1) \textit{happy} single-task dialogs, (2) \textit{unhappy} single-task dialogs and (3) \textit{multi-task} dialogs. Here, \textit{happy} refers to dialogs where the users are cooperative and complete the task. In contrast, \textit{unhappy} dialogs consist of uncooperative users that may change the subject, engage in irrelevant chit-chat and otherwise aim to push the system beyond its capabilities. Since our primary objective is to address zero-shot transfer, we only consider the \textit{happy} single-task dialogs. There are 1537 happy single-task dialogs and 10,364 turns.

\subsection{Zero-Shot Setting}

In the STAR dataset, there are 23 dialog tasks (13 domains) with \textit{happy} single-task dialogs. We perform two types of transfer learning experiments: task transfer and domain transfer. In task transfer, a model is trained on $n-1$ tasks (i.e., 22) and evaluated on the last one. This is repeated for each of the 23 tasks. For domain transfer, a model is trained on $n-1$ domains (i.e., 12) and evaluated on the last one. In task transfer, there may be some overlap between the training and testing, for example, the domain-specific terminology. In contrast, in domain transfer there is very limited overlap. When the model is tasked with generalizing to the restaurant domain, it has seen nothing related to restaurants during training.

In both of these settings, the model is aware of which task it is being evaluated on, meaning that it can leverage a task specification (e.g., schema) for the new task. This experimental design resembles a real-world setting where a system developer would be aware of the new task. For example, if a developer wanted to extend a dialog system to handle a COVID-19 related question, they would be able to create a new task specification. As such, our goal is to develop a model that can generalize to an unseen task conditioned on a task specification.

\section{Methods}

In order to enable zero-shot transfer to new dialog tasks and domains, the Schema Attention Model (\textsc{sam}) is introduced. It leverages an external dialog policy representation (i.e., the schema) to predict the next system action. This section begins by describing the baseline model for the task of next action prediction. Next, the schema-guided paradigm is introduced (Figure \ref{fig:schema_paradigm}). It includes a graph-based representation of the task-specific schema and \textsc{sam}, a model that identifies the next system action by attending to a task-specific schema representation.

\subsection{Baseline}

This section describes the baseline model proposed by \citet{mosig2020star}. Given an arbitrary language encoder, denoted as $\mathcal{F}$, the baseline model obtains a vector representation of the dialog history, $c$. This representation is then passed through a softmax layer to obtain a probability distribution over the actions.

\begin{gather}
    {\bm{h}} = \mathcal{F}( c ) \\
    P_{\text{clf}} = \text{softmax} ( {\bm{W}} {\bm{h}}^T + {\bm{b}} )
\end{gather}

Throughout this paper, BERT-base \citep{devlin-etal-2019-bert} is used as the language encoder.

\subsection{Schema-Guided Paradigm}

Our baseline model simultaneously needs to (1) interpret the dialog context and identify the relevant intents and slots, and (2) learn the task-specific dialog policies (i.e., if the user wants the weather, ask the city) for the different tasks in the training data. This model is incapable of generalizing to a new task in a zero-shot setting, as it would lack knowledge of the task-specific policy for the new task. To mitigate this problem and to enable zero-shot task transfer, we present the schema-guided paradigm which decouples the task-specific dialog policy from the language understanding.

An example is shown in Figure \ref{fig:schema_paradigm}: the schema-guided paradigm decouples the the dialog policy from language understanding by explicitly providing task-specific schema graphs as input to the model. These schema graphs serve as complete representations of the dialog policy for a given task. Therefore, while the baseline needs to implicitly learn the dialog policies, a schema-guided model instead learns to leverage the explicit schema graphs. As such, a schema-guided model can generalize to a new task as long as it is provided with the corresponding schema graph.

In this paradigm, the role of the model is to interpret a dialog context and align it to the explicit schema graph. The role of the schema graph is to determine the next action according to the dialog policy. In this manner, the language encoder is being trained for the task of sentence similarity. With the help of pre-trained models, language understanding in a schema-guided paradigm can be considered to be task-agnostic. By decoupling the task-agnostic language understanding and the task-specific dialog policy, the schema-guided paradigm better facilitates zero-shot transfer learning.

The schema-guided paradigm consists of the representation of the schema graph, and a neural model which interprets the dialog context and aligns it to the schema graph.

\subsubsection{Schema Representation}

In the schema-guided paradigm, the schema representation is the task-specific dialog policy. To ensure the efficacy and robustness of the  dialog system, it is important that the schema representation be complete and informative. In the case of ambiguity or incompleteness in the schema representation, the next action will fail to be correctly predicted, regardless of the strength of the model. The schema representations are manually constructed for every task. In the schema-guided paradigm, to transfer to a new task, a system developer would simply need to construct a new schema representation.

\citet{mosig2020star} propose a baseline schema representation wherein the nodes of the graph correspond to system actions and database states. There are nodes for user states only in situations where the system behavior differs depending on the user's actions (e.g., \textit{`Yes'} $\rightarrow$ \textit{ask-time}, \textit{`No'} $\rightarrow$ \textit{ask-date}). The consequence of this representation is that when the model aligns the dialog history to the schema, it largely relies on the system utterances. However, this representation fails to account for realistic user behavior and therefore yields only marginal improvement over the baseline.

Specifically, users will often provide information out of turn (e.g., \textit{System: `Where would you like to go?' $\rightarrow$ User: `Leaving from the airport and going downtown'}). In this example, it is difficult for the model to realize that the question \textit{System: `Where are you leaving from?'} has also been answered and therefore should not be the next system action. Users can also ignore the system utterance (e.g., \textit{System: `Where would you like to go?'} $\rightarrow$ \textit{User: `Actually, what's the weather?'}). It is thus ineffective to represent dialog policy only in terms of the system utterances. To this end, we extend the schema representation by incorporating user utterances into the schema graph.

As shown in Figure \ref{fig:schema_graph}, our schema graph incorporates nodes corresponding to user utterances. As such, if a user provides information out of turn or changes the subject, our model will be able to effectively align the dialog to the schema. To account for variance in the user utterances, future work could extend this schema representation to include multiple variations of a given user utterance. However, as the schema graphs are manually constructed for every task, there is a trade-off between manual effort and efficacy\footnote{Constructing the schema graphs is not particularly labor-intensive. It took the first author between 15 and 45 minutes to create each schema graph, depending on the complexity of the task.}.

\begin{figure}
    \centering
    \includegraphics[width=0.5\textwidth]{example_schema.png}
    \caption{A section of the task-specific schema graph for the \textit{bank-balance} task. The system must authenticate the user with their account number and PIN. However, if the user has forgotten either of these, it must ask backup security questions. The blue nodes correspond to system actions and the yellow nodes denote user utterances.}
    \label{fig:schema_graph}
\end{figure}

The schema graph has several noteworthy properties. First, the system actions are consistently deterministic. Nodes corresponding to a database response or to a user utterance will always have a single outgoing edge to a system response node. Furthermore, such nodes will also have a single incoming edge from a system response node. For a given user/database node, $u$, we denote the previous system response node as $\textbf{prev}(u)$ and the following system response as $\textbf{next}(u)$. Each node has some text associated with it, denoted as $\textbf{text}(u)$. This text is a template for either a system utterance, database response or user utterance. System nodes will also have an associated system action, $\textbf{act}(u)$. There is a one-to-one mapping between the system actions and the system response templates.

\subsubsection{Schema Attention Model}

In the schema-guided paradigm, the role of the model is to understand the dialog history and align it to the schema representation. We introduce the \textbf{Schema Attention Model}, \textsc{sam}, which attends between the dialog history, $c = {c_1, \dots, c_N}$ and the schema graph. \textsc{sam} extends the schema-guided model presented by \citet{mosig2020star} by (1) leveraging a stronger attention mechanism, (2) improving the training algorithm, and (3) removing the linear classification layer which is detrimental to zero-shot performance.

The objective of \textsc{sam} is to predict the node in the schema graph that best corresponds to the dialog context. \textsc{sam} will produce a probability distribution over the nodes corresponding to user utterances and database responses. Given an attention distribution over the nodes, we can obtain a probability distribution over the set of actions by propagating the attention probabilities over the graph. Concretely, if node $u$ has an attention weight of $p$, we add $p$ to the probability of $\textbf{action}(\textbf{next}(u))$.

We consider every node $u$ that corresponds to either a database response or a user utterance. We then represent each node $u$ as the concatenation of the previous node and the current node, i.e., $\textbf{text}(\textbf{prev}(u)) + \textbf{text}(u)$. For all nodes $u \in U$, we obtain this textual representation denoted as $s \in S$.

We are given a language encoder, $\mathcal{F}$, the dialog context, $c = {c_1, \dots, c_N}$, the nodes $U$, their corresponding textual representations $S$, and the set of possible actions $A$. Note that unlike in Equation 1, $\mathcal{F}$ is used to produce a vector representation of each word in the input. \textsc{sam} produces a probability distribution over the actions as follows:

\begin{gather}
    {\bm{h}}_{1,\dots,N} = \mathcal{F}(c \colon c_1, \dots, c_N) \\
    {\bm{S}}_{i ; 1,\dots,M} = \mathcal{F}(S_i \colon s_1,\dots,s_M) \\
    {\bm{w}}^{i}_{j,k} = {\bm{h}}_j^T  {\bm{S}}_{i;k} \\
    \mathbf{\alpha} = \text{softmax}({\bm{w}}^{1,\dots,|S|})\\
    p_i = \sum_{j \leq N} \sum_{k \leq M} \alpha^i_{j,k}
\end{gather}

Here, ${\bm{w}}^i$ is an $N \times M$ dimensional matrix corresponding to the dot product between the $N$ words of the dialog history and the $M$ words of the $i$-th textual representation in $S$. To get the attention weights over all of the words of the schema, we perform a softmax over all ${\bm{w}}^i, 1 \leq i \leq |S|$. By summing over the attention weights in $\alpha^i$, we get $p_i$, a scalar value which denotes the attention between the dialog history and the $i$-th node (i.e., the corresponding textual representation $S_i$). Given $p_i$ we produce a probability distribution over the actions $A$ as follows:

\begin{gather}
    g(i,a) = \begin{cases}
        p_i, & \text{if } \textbf{action}(\textbf{next}(u_i)) = a\\
        0, & \text{otherwise}
    \end{cases}\\
    P(a) = \sum_{i \leq |S|} g(i,a)
\end{gather}

To align the dialog history to the schema graph, \textsc{sam} performs word-level attention using a BERT-base model. In contrast, the schema-guided model of \citet{mosig2020star} attends with the sentence level vector representation produced by BERT. With the word-level attention, \textsc{sam} can better align ambiguous dialog contexts, such as situations where the user provides multiple pieces of information in a single utterance. Since this word-level attention operates on the sub-word tokens used in BERT, it can also potentially handle spelling errors in the user utterances.

Furthermore, in their schema-guided model, \citet{mosig2020star} combine the probability distribution produced by attending to the schema graph with their baseline model (i.e., Section 3.1). While this may result in better performance on the tasks the model is trained with, the baseline model will not generalize to unseen tasks. In contrast, \textsc{sam} computes the probability for an action using only the attention over the schema graph.

\citet{mosig2020star} train their schema-guided model to predict the appropriate node, $u_i$, from a set of nodes $U'$ (s.t., $U' \subset U$). At training time, for efficiency reasons, the set of nodes $U'$ is obtained by using the corresponding node for every dialog context in the training batch. Since the training batches are randomly sampled, this results in $U'$ including nodes from a variety of different schema graphs. At inference time, the dialog task is known and therefore only the corresponding schema graph needs to be attended to (i.e., $U'$ will contain nodes from a single schema graph). It is valuable to train the model to distinguish between different nodes of the same schema graph. Specifically, the attention mechanism (i.e., Equations 5 - 6) will learn stronger fine-grained relationships when trained with negative samples from the \textit{same} domain. As such, we augment the training algorithm to sample batches from the same dialog task, meaning that $U'$ will only include nodes from a single schema.

\textsc{sam} improves on the baseline schema-guided model introduced by \citet{mosig2020star} by (1) leveraging a stronger attention mechanism that better handles realistic user behavior, (2) computing a probability distribution \textit{only} by attending to the schema graph and (3) modifying the training algorithm to have in-domain negative samples which result in the model learning to identify fine-grained relationships. In combination with the improved schema representation, \textsc{sam} is better suited to handle realistic user behavior in zero-shot settings.
\begin{table}[]
\renewcommand*{\arraystretch}{1.1}

    \centering
{%
\begin{tabularx}{0.4\textwidth}{@{\extracolsep{\stretch{1}}} c *{2}{c}}
\toprule
\textbf{Model} & \textbf{$F_1$ score}  & \textbf{Accuracy}  \\ \midrule
Baseline $\diamond$  & \textbf{73.79} & \textbf{74.85}  \\
\textsc{bert+s} $\diamond$ & 71.59 & 72.27 \\
\textsc{sam} $-$ [1] & 54.35 & 60.51 \\
\textsc{sam} $-$ [2,3,4] & 70.22 & 71.01 \\
\textsc{sam} $-$ [2] & 70.27 & 71.93  \\
\textsc{sam} $-$ [3] & 70.18 & 71.64 \\
\textsc{sam} $-$ [4] & 69.68 & 69.79 \\
\textsc{sam} & 70.38 & 71.45 \\
\bottomrule
\end{tabularx}
}
    \caption{Performance in the standard experimental setting. Models marked with $\diamond$ are attributed to \citet{mosig2020star}. We denote their schema-guided model, `BERT + \textit{Schema}', as \textit{\textsc{bert+s}}. \textsc{sam} consists of four improvements upon \textsc{bert+s}: (1) user-aware schema, (2) word-level attention, (3) using negative samples from the same task at training, (4) removing the linear classification layer. Results in boldface are statistically significant by t-test ($p < 0.01$)}
    \label{tab:standard}
\end{table}

\begin{table*}[]
\renewcommand*{\arraystretch}{1.1}

    \centering
{%
\begin{tabularx}{0.8\textwidth}{@{\extracolsep{\stretch{1}}} c *{4}{c}}
\toprule
\textbf{Model} & \multicolumn{2}{c}{\textbf{Task Transfer}} & \multicolumn{2}{c}{\textbf{Domain Transfer}}  \\
& \textbf{$F_1$ score} & \textbf{Accuracy} & \textbf{$F_1$ score} & \textbf{Accuracy} \\
\midrule
Baseline $\diamondsuit$  & 31.23 & 30.65 & 31.82 & 33.92 \\
\textsc{bert+s} $\diamondsuit$ & 28.12 & 28.28 & 29.70  & 32.43\\
\textsc{sam} $-$ [1] & 33.81 & 37.84 & 41.77 & 45.64  \\
\textsc{sam} $-$ [2,3,4] & 43.28 & 46.11 & 43.78 & 45.19 \\
\textsc{sam} $-$ [2] & 50.72 & 53.69 & 52.20  & 54.68 \\
\textsc{sam} $-$ [3] & 45.54 & 49.29 & 50.56 & 52.13 \\
\textsc{sam} $-$ [4] & 47.26 & 47.99 & 47.67 & 48.92 \\
\textsc{sam} & \textbf{53.31} & \textbf{55.51} & \textbf{55.74} & \textbf{57.75}\\
\bottomrule
\end{tabularx}
}
    \caption{Performance in zero-shot transfer. We present results on both task transfer and domain transfer. Models marked with $\diamondsuit$ are attributed to \citet{mosig2020star}. \textsc{sam} consists of four improvements upon \textsc{bert+s}: (1) user-aware schema, (2) word-level attention, (3) using negative samples from the same task at training, (4) removing the linear classification layer. Results in bold-face are statistically significant by t-test ($p < 0.01$).}
    \label{tab:zero}
\end{table*}

\section{Results}

To validate the effectiveness of \textsc{sam}, a number of \textit{next action prediction} experiments are carried out on the STAR dataset \citep{mosig2020star}. First, \textsc{sam} is evaluated in the standard experimental setting, i.e., training and testing on the same tasks. Next, we carry out zero-shot transfer experiments, as defined in Section 3. The evaluation uses accuracy and weighted $F_1$ score.

We rerun the experiments presented by \citet{mosig2020star} using code shared by the authors. In our results, the model introduced by \citet{mosig2020star} is denoted as \textsc{bert+s}. Their original results were obtained on an older version of STAR, with annotation errors\footnote{Specifically, certain dialogs were misattributed as being \textit{happy} single-task dialogs.} that have since been fixed.

\subsection{Standard Experiments}

In the standard experimental setting, models are trained and tested on the same tasks. Following \citet{mosig2020star}, 80\% of the dialogs are used for training and 20\% for testing. All models are trained for 50 epochs.

The results shown in Table \ref{tab:standard} show \textsc{sam} to be comparable to the baseline model on the standard setting. Since the augmentations to \textsc{sam} are primarily intended to improve zero-shot performance, it is unsurprising that there is no performance improvement compared to the standard setting. When evaluating on \textit{seen} tasks, the linear classification layer is significantly more effective than attending to the schema. This suggests that a large neural model (i.e., BERT) is able to implicitly learn meaningful dialog policies from dialog data. It is possible that this performance difference may decrease with more expressive schemas (e.g., having multiple examples for each user utterance, automatically learning schemas from the dataset). The value of our schema graphs is nonetheless shown when comparing \textsc{sam} to \textsc{sam}$-[1]$ (i.e., the old schema graphs). These experiments provide an upper bound for the performance in zero-shot transfer.

\subsection{Zero-Shot Transfer}

Table \ref{tab:zero} shows the results of the zero-shot experiments. \textsc{sam} obtains strong improvements over the baseline models for both zero-shot task transfer and domain transfer. These experimental results validate the effectiveness of the schema-guided paradigm, as well as the specific design of \textsc{sam}.

Compared to the baseline model (described in Section 3.1), \textsc{sam} obtains a $\mathbf{+22 ~F_1}$ score improvement in task transfer and a $\mathbf{+24 ~F_1}$ score improvement in domain transfer. Since the baseline model is unable to predict classes it has not observed at training time, its performance is limited to actions that are consistent across domains (e.g., \textit{`hello'}, \textit{`goodbye'}, \textit{`anything-else'}). This improvement highlights the effectiveness of the schema-guided paradigm for zero-shot transfer learning.

\textsc{bert+s} also leverages schemas for transfer learning. Yet, it under-performs relative to the baseline model. \textsc{sam} attains even larger improvements over this baseline schema-guided model. As described in Section 4.2, the weak performance of \textsc{bert+s} is largely a consequence of it being incapable of handling realistic user behavior. The design of \textsc{bert+s} (i.e., the schema only having system nodes) results in the model essentially predicting the subsequent system actions. This is equivalent to sequentially predicting the next system action, regardless of user behavior. With improved schema representations and model architecture, \textsc{sam} achieves much stronger performance in zero-shot transfer.

Our ablation experiments shed more light on the performance of \textsc{sam} relative to \textsc{bert+s}. A significant performance drop is observed when removing the newly constructed schema representations (i.e., \textsc{sam}$- [1]$). In contrast, adding the schema graphs to \textsc{bert+s} (i.e., \textsc{sam}$- [2,3,4]$) results in a strong performance improvement of $\mathbf{+15~F_1}$ score. This confirms the hypothesis that the schema graphs of \citet{mosig2020star}, which are largely comprised of system action nodes are insufficient for modelling realistic user behavior.

Word-level attention is shown to give moderate, albeit statistically significant, improvement. In contrast to \textsc{sam}$-[2]$, \textsc{sam} obtains a $\mathbf{+3~F_1}$ score improvement. While word-level attention allows the model to better align the dialog to the schema, it is an architectural improvement that is not central to the schema-guided paradigm.

Modifying the training algorithm to sample batches from the same task results in better negative samples during training. This allows the model to learn to distinguish between nodes from the same schema graph when aligning the dialog to the schema graph. When this modification is removed (i.e., \textsc{sam}$-[3])$, the performance of \textsc{sam} drops by $\mathbf{8~F_1}$ score for zero-shot task transfer.

The fourth and final component of \textsc{sam} is the removal of the linear classification layer. Since this classification layer is unable to predict classes it has not seen at training time, it is ineffective in zero-shot settings. Unsurprisingly, removing it increases performance and \textsc{sam} obtains a $\mathbf{+6~F_1}$ score improvement over \textsc{sam}$-[4]$.

The zero-shot experiments shown in Table \ref{tab:zero} empirically validate several hypotheses presented in this paper. First, the strong improvement over the baseline demonstrates the efficacy of the schema-guided paradigm for zero-shot generalizability in end-to-end dialog. Decoupling dialog policy and the language understanding by explicitly representing the task-specific dialog policies as schema graphs results in an improved ability to transfer to unseen tasks. Next, we improve over the schema-guided model of \citet{mosig2020star} through (1) an improved schema representation and (2) a collection of modifications to the model. The improved schema representation better models realistic user behaviors in dialog, and therefore results in better alignment of the dialog and the schema. Our model modifications result in the model being able to learn better fine-grained relationships during alignment (e.g., through better negative sampling and word-level attention) and better handle zero-shot transfer (e.g., by removing the linear layer).

In contrast to prior work on zero-shot generalizability \citep{zhao2018zero,qian2019domain}, our approach is shown to effectively  transfer between the vastly dissimilar domains of the STAR corpus \citep{mosig2020star} (e.g., trivia or spaceship maintenance). Rather than modelling a cross-domain mapping and leveraging similar concepts across different domains, the schema-guided paradigm \textit{decouples} the domain-specific (i.e., the dialog policy) and domain-agnostic (i.e., language understanding) aspects of dialog systems. Through the schema-guided paradigm, we achieve strong performance in the zero-shot setting and take an important step towards zero-shot dialog.

\section{Conclusion}

This paper shows strong results in zero-shot task transfer and domain transfer using the schema-guided paradigm. We hypothesized that the difficulty of zero-shot transfer in dialog stems from the dialog policy. When neural models implicitly memorize dialog policies observed at training time, they struggle to transfer to new tasks. To mitigate this, we explicitly provide the dialog policy to the model, in the form of a schema graph. This paper introduces the Schema Attention Model (\textsc{sam}) and shows improved schema graphs for the STAR corpus. This approach attains significant improvement over prior work in the zero-shot setting, with a \textbf{$\mathbf{+22 ~F_1}$ score improvement}. Furthermore, the ablation experiments demonstrate the effectiveness of both \textsc{sam} and the improved schema representations. Future work may explore (1) improved schema representations to better capture dialog policy, (2) improved model architectures to better align the dialog to the schema, and (3) extensions to other problems (e.g., response generation).
\bibliographystyle{acl_natbib}