A search for direct top squark pair production resulting in events with either a same-flavour opposite-sign dilepton pair with invariant mass compatible with a Z boson or a pair of jets compatible with a Standard Model (SM) Higgs boson (h) is presented. Requirements on the missing transverse momentum, together with additional selections on leptons, jets, jets identified as originating from b-quarks are imposed to target the other decay products of the top squark pair. The analysis is performed using proton-proton collision data at √s=13 TeV collected with the ATLAS detector at the LHC in 2015–2016, corresponding to an integrated luminosity of 36.1 fb−1. No excess is observed in the data with respect to the SM predictions. The results are interpreted in two sets of models. In the first set, direct production of pairs of lighter top squarks (t~1) with long decay chains involving Z or Higgs bosons is considered. The second set includes direct pair production of the heavier top squark pairs (t~2) decaying via t~2→Zt~1 or t~2→ht~1. The results exclude at 95% confidence level t~2 and t~1 masses up to about 800 GeV, extending the exclusion region of supersymmetric parameter space covered by previous LHC searches.